[Refactor] Move weight completeness verification into the loader

python/sglang/srt/model_executor/model_runner_components/weight_updater.py

@@ -178,7 +178,18 @@ def get_weight_iter(config):
             return iter
 
         def model_load_weights(model, iter):
-            loader.load_weights_and_postprocess(model, iter, target_device)
+            # Elastic-EP refills may select only expert tensors, while FLASH_RL
+            # refits may arrive in batches even with another disk load format.
+            # Neither establishes model-wide checkpoint completeness.
+            is_full_checkpoint = weight_name_filter is None and not getattr(
+                model, "flash_rl_initial_load_complete", False
+            )
+            loader.load_weights_and_postprocess(
+                model,
+                iter,
+                target_device,
+                is_full_checkpoint=is_full_checkpoint,
+            )
             return model
 
         with set_default_torch_dtype(self.model_config.dtype):

python/sglang/srt/model_loader/loader.py

@@ -86,6 +86,9 @@
     get_model_architecture,
     set_default_torch_dtype,
 )
+from sglang.srt.model_loader.weight_completeness import (
+    load_and_verify_weights as _load_and_verify_weights,
+)
 from sglang.srt.utils.common import is_cuda_alike
 
 # Constants for memory management
@@ -804,7 +807,9 @@ def load_model(
         return model.eval()
 
     @staticmethod
-    def load_weights_and_postprocess(model, weights, target_device):
+    def load_weights_and_postprocess(
+        model, weights, target_device, *, is_full_checkpoint: bool = True
+    ):
         # Used in tests to verify memory savings when using online quantization.
         if is_cuda_alike():
             peak_memory = torch.cuda.max_memory_allocated()
@@ -823,12 +828,16 @@ def load_weights_and_postprocess(model, weights, target_device):
             # Scope exact FP4 quantization math to load-time conversion only;
             # restore the original environment before serving starts.
             with temp_set_env(FLASHINFER_DISABLE_FP4_QUANT_FAST_MATH="1"):
-                model.load_weights(weights)
+                _load_and_verify_weights(
+                    model, weights, is_full_checkpoint=is_full_checkpoint
+                )
             if target_device.type == "cuda":
                 torch.cuda.synchronize()
                 torch.cuda.empty_cache()
         else:
-            model.load_weights(weights)
+            _load_and_verify_weights(
+                model, weights, is_full_checkpoint=is_full_checkpoint
+            )
 
         # Used in tests to verify memory savings when using online quantization.
         if is_cuda_alike():
@@ -1005,13 +1014,21 @@ def _bind_method_to_cls(func, obj):
             return func
         return types.MethodType(func, obj)
 
-    def load_weights_and_postprocess(self, model, weights, target_device):
+    def load_weights_and_postprocess(
+        self, model, weights, target_device, *, is_full_checkpoint: bool = True
+    ):
         """
         Initial load: Load BF16 → Record state → Apply FP8 quantization.
-        Called ONCE during model initialization.
+        Online reload: Reuse the installed fast-path proxy.
         """
-        logger.info("[QuantizedRL] Initial load with FP8 quantization")
+        is_reload = QuantizedRLModelLoader.is_reload_scenario(model)
+        if is_reload:
+            # Reloads may arrive as several batches, even when a disk iterator
+            # is unfiltered, so they do not establish checkpoint completeness.
+            model.load_weights(weights)
+            return
 
+        logger.info("[QuantizedRL] Initial load with FP8 quantization")
         original_load_weights = model.load_weights
 
         def load_weights_proxy(weights):
@@ -1021,11 +1038,11 @@ def load_weights_proxy(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, is_full_checkpoint=is_full_checkpoint)
         original_weights = dict(model.named_parameters())
 
         # Record pre-quantization state (shape/stride) for torch.as_strided reset
@@ -2856,7 +2873,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()
 
@@ -3048,8 +3065,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():
@@ -3363,7 +3381,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)
@@ -3433,7 +3451,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,43 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from collections.abc import Callable, Container, Iterable
+
+
+def unloaded_required_params(
+    parameter_names: Iterable[str],
+    loaded: Container[str],
+    is_optional: Callable[[str], bool],
+) -> set[str]:
+    """Return required model parameters that a checkpoint did not load."""
+    return {
+        name for name in parameter_names if name not in loaded and not is_optional(name)
+    }
+
+
+def load_and_verify_weights(
+    model, weights, *, is_full_checkpoint: bool = True
+) -> object:
+    """Load weights and, for migrated models, verify a full checkpoint."""
+    loaded_params = model.load_weights(weights)
+    if not is_full_checkpoint or not getattr(model, "verify_weights_on_load", False):
+        return loaded_params
+
+    if loaded_params is None:
+        raise TypeError(
+            f"{type(model).__name__}.load_weights() must return loaded parameter "
+            "names when verify_weights_on_load is enabled"
+        )
+
+    is_optional = getattr(model, "is_optional_weight", lambda _name: False)
+    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

python/sglang/srt/models/deepseek_ocr.py

@@ -1423,6 +1423,8 @@ def build_qwen2_decoder_as_encoder(
 
 
 class DeepseekOCRForCausalLM(nn.Module):
+    verify_weights_on_load = True
+
     def __init__(
         self,
         *,
@@ -1855,12 +1857,8 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
                 weight_loader = getattr(param, "weight_loader", default_weight_loader)
                 weight_loader(param, loaded_weight)
             loaded_params.add(name)
-        unloaded_params = params_dict.keys() - loaded_params
-        if unloaded_params:
-            raise RuntimeError(
-                f"Some weights are not initialized from checkpoints: {unloaded_params}"
-            )
         self.post_load_weights()
+        return loaded_params
 
     def post_load_weights(self):
         if _is_cpu and _is_cpu_amx_available:

python/sglang/srt/models/deepseek_v4.py

@@ -2459,6 +2459,8 @@ def forward(
 
 
 class DeepseekV4ForCausalLM(nn.Module):
+    verify_weights_on_load = True
+
     def __init__(
         self,
         config: DeepSeekV4Config,
@@ -3114,39 +3116,27 @@ def auto_weight_loader(module):
             for future in concurrent.futures.as_completed(futures):
                 future.result()
 
+        # A filtered reload may select a subset of runtime parameters, but each
+        # fused target still requires all of its checkpoint shards.
         assert len(cache_compressor_weight) == 0
         assert len(cache_wqkv_a_weight) == 0, cache_wqkv_a_weight.keys()
-        unloaded_params = params_dict.keys() - loaded_params
+        self.post_load_weights(is_nextn=is_nextn, weight_names=weight_names)
 
-        skipped_checking_patterns = [
+        if not is_nextn:
+            self._prewarm_mhc_pre_kernels()
+        return loaded_params
+
+    def is_optional_weight(self, name: str) -> bool:
+        patterns = [
             "attn_mqa.k_scale",
             "attn_mqa.v_scale",
             "blockscale_swizzled",
         ]
         if not self.pp_group.is_first_rank:
-            skipped_checking_patterns.append("embed_tokens")
+            patterns.append("embed_tokens")
         if not self.pp_group.is_last_rank:
-            skipped_checking_patterns.append("model.norm.")
-            skipped_checking_patterns.extend(["lm_head", "hc_head_"])
-        if is_nextn:
-            skipped_checking_patterns.extend(["lm_head", "embed_tokens"])
-        unloaded_params = {
-            p
-            for p in unloaded_params
-            if all(
-                skipped_checking_pattern not in p
-                for skipped_checking_pattern in skipped_checking_patterns
-            )
-        }
-        if unloaded_params:
-            logger.warning(
-                f"Some weights are not initialized from checkpoints: {unloaded_params}"
-            )
-
-        self.post_load_weights(is_nextn=is_nextn, weight_names=weight_names)
-
-        if not is_nextn:
-            self._prewarm_mhc_pre_kernels()
+            patterns.extend(["model.norm.", "lm_head", "hc_head_"])
+        return any(pattern in name for pattern in patterns)
 
     def get_embed_and_head(self):
         return self.model.embed_tokens.weight, self.lm_head.weight

python/sglang/srt/models/deepseek_v4_nextn.py

@@ -274,7 +274,12 @@ def forward(
         )
 
     def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
-        super().load_weights(weights, is_nextn=True)
+        return super().load_weights(weights, is_nextn=True)
+
+    def is_optional_weight(self, name: str) -> bool:
+        if "lm_head" in name or "embed_tokens" in name:
+            return True
+        return super().is_optional_weight(name)
 
     def post_load_weights(self, is_nextn=False, weight_names=None):
         super().post_load_weights(is_nextn=True, weight_names=weight_names)

python/sglang/srt/models/gemma4_causal.py

@@ -12,7 +12,6 @@
 # limitations under the License.
 # ==============================================================================
 
-import logging
 import re
 from typing import Iterable, List, Optional, Set, Tuple, Union
 
@@ -61,8 +60,6 @@
 from sglang.srt.runtime_context import get_parallel, get_server_args
 from sglang.srt.utils import add_prefix, make_layers
 
-logger = logging.getLogger(__name__)
-
 
 # Aligned with HF's implementation, using sliding window inclusive with the last token
 # SGLang assumes exclusive
@@ -1024,6 +1021,7 @@ def forward(
 
 
 class Gemma4ForCausalLM(PreTrainedModel):
+    verify_weights_on_load = True
     config_class = Gemma4TextConfig
     base_model_prefix = "language_model"
     _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
@@ -1218,12 +1216,6 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
 
         params_dict = dict(self.named_parameters())
         params_dict.update(dict(self.named_buffers()))
-        non_persistent_buffers: Set[str] = set()
-        for mod_name, mod in self.named_modules():
-            for buf_name in getattr(mod, "_non_persistent_buffers_set", set()):
-                full = f"{mod_name}.{buf_name}" if mod_name else buf_name
-                non_persistent_buffers.add(full)
-
         loaded_params: Set[str] = set()
         for name, loaded_weight in weights:
             name = name.replace("model.language_model.", "model.")
@@ -1346,27 +1338,6 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
                         )
                         weight_loader(param, loaded_weight)
                         loaded_params.add(name)
-        unloaded_params = params_dict.keys() - loaded_params
-        if unloaded_params:
-            param_names = set(dict(self.named_parameters()).keys())
-            buckets = {
-                logging.WARNING: (
-                    "Some weights are not initialized from checkpoints",
-                    lambda p: p in param_names,
-                ),
-                logging.INFO: (
-                    "Persistent buffers not in checkpoint (using default init)",
-                    lambda p: p not in param_names and p not in non_persistent_buffers,
-                ),
-                logging.DEBUG: (
-                    "Non-persistent buffers not in checkpoint (expected)",
-                    lambda p: p in non_persistent_buffers,
-                ),
-            }
-            for level, (msg, pred) in buckets.items():
-                names = sorted(p for p in unloaded_params if pred(p))
-                if names:
-                    logger.log(level, "%s: %s", msg, names)
         return loaded_params
 
     def _shard_weight(self, weight: torch.Tensor) -> torch.Tensor:

python/sglang/srt/models/gemma4_mm.py

@@ -128,6 +128,7 @@ def forward(
 
 
 class Gemma4ForConditionalGeneration(PreTrainedModel):
+    verify_weights_on_load = True
     config_class = Gemma4Config
     """Gemma4 multimodal model for conditional generation."""
 
@@ -860,12 +861,6 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
 
         params_dict = dict(self.named_parameters())
         params_dict.update(dict(self.named_buffers()))
-        non_persistent_buffers: Set[str] = set()
-        for mod_name, mod in self.named_modules():
-            for buf_name in getattr(mod, "_non_persistent_buffers_set", set()):
-                full = f"{mod_name}.{buf_name}" if mod_name else buf_name
-                non_persistent_buffers.add(full)
-
         text_tie = getattr(self.config.text_config, "tie_word_embeddings", True)
         start_layer = self.language_model.start_layer
         end_layer = self.language_model.end_layer
@@ -1017,27 +1012,6 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
                         )
                         weight_loader(param, loaded_weight)
                         loaded_params.add(name)
-        unloaded_params = params_dict.keys() - loaded_params
-        if unloaded_params:
-            param_names = set(dict(self.named_parameters()).keys())
-            buckets = {
-                logging.WARNING: (
-                    "Some weights are not initialized from checkpoints",
-                    lambda p: p in param_names,
-                ),
-                logging.INFO: (
-                    "Persistent buffers not in checkpoint (using default init)",
-                    lambda p: p not in param_names and p not in non_persistent_buffers,
-                ),
-                logging.DEBUG: (
-                    "Non-persistent buffers not in checkpoint (expected)",
-                    lambda p: p in non_persistent_buffers,
-                ),
-            }
-            for level, (msg, pred) in buckets.items():
-                names = sorted(p for p in unloaded_params if pred(p))
-                if names:
-                    logger.log(level, "%s: %s", msg, names)
         return loaded_params
 
     lora_pattern = re.compile(