From 62d741851005361f420b4beb19ac51853efa39a3 Mon Sep 17 00:00:00 2001 From: EazyReal <8047065+EazyReal@users.noreply.github.com> Date: Thu, 30 Jul 2026 05:42:54 +0000 Subject: [PATCH] model_loader: verify full-checkpoint weight completeness --- .../model_runner_components/weight_updater.py | 13 ++- python/sglang/srt/model_loader/loader.py | 44 +++++++--- .../srt/model_loader/weight_completeness.py | 43 +++++++++ python/sglang/srt/models/deepseek_ocr.py | 8 +- python/sglang/srt/models/deepseek_v4.py | 38 +++----- python/sglang/srt/models/deepseek_v4_nextn.py | 7 +- python/sglang/srt/models/gemma4_causal.py | 31 +------ python/sglang/srt/models/gemma4_mm.py | 28 +----- python/sglang/srt/models/gemma4_unified.py | 30 ------- python/sglang/srt/models/mllama4.py | 8 +- .../test_load_and_verify_weights.py | 87 +++++++++++++++++++ 11 files changed, 201 insertions(+), 136 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_executor/model_runner_components/weight_updater.py b/python/sglang/srt/model_executor/model_runner_components/weight_updater.py index ecbfa7a713f7..6c6e16677d4e 100644 --- a/python/sglang/srt/model_executor/model_runner_components/weight_updater.py +++ b/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): diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index 8bafdd5cd6f1..fb75e29ba780 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/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) 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..d2942e3d5c76 --- /dev/null +++ b/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 diff --git a/python/sglang/srt/models/deepseek_ocr.py b/python/sglang/srt/models/deepseek_ocr.py index 1b70614f6db8..d39c7ba53e9a 100644 --- a/python/sglang/srt/models/deepseek_ocr.py +++ b/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: diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index a21936d7352a..096eda773b50 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/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 diff --git a/python/sglang/srt/models/deepseek_v4_nextn.py b/python/sglang/srt/models/deepseek_v4_nextn.py index 1dd326c6e717..b54ac05681f4 100644 --- a/python/sglang/srt/models/deepseek_v4_nextn.py +++ b/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) diff --git a/python/sglang/srt/models/gemma4_causal.py b/python/sglang/srt/models/gemma4_causal.py index f2d9a645e6c2..7d040aab2ac3 100644 --- a/python/sglang/srt/models/gemma4_causal.py +++ b/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: diff --git a/python/sglang/srt/models/gemma4_mm.py b/python/sglang/srt/models/gemma4_mm.py index 56d033ff4b63..a30febfbae86 100644 --- a/python/sglang/srt/models/gemma4_mm.py +++ b/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( diff --git a/python/sglang/srt/models/gemma4_unified.py b/python/sglang/srt/models/gemma4_unified.py index 67d0771d4aea..a00600c67c91 100644 --- a/python/sglang/srt/models/gemma4_unified.py +++ b/python/sglang/srt/models/gemma4_unified.py @@ -32,7 +32,6 @@ only construction, per-modality feature extraction and weight loading. """ -import logging import re from typing import Iterable, List, Optional, Set, Tuple @@ -55,8 +54,6 @@ from sglang.srt.models.gemma4_mm import Gemma4ForConditionalGeneration from sglang.srt.utils import add_prefix -logger = logging.getLogger(__name__) - class Gemma4UnifiedVisionEmbedder(nn.Module): """Encoder-free vision embedder. @@ -342,12 +339,6 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> Set[str]: 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 @@ -411,27 +402,6 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> Set[str]: 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 diff --git a/python/sglang/srt/models/mllama4.py b/python/sglang/srt/models/mllama4.py index 68fc8f6db046..d21a50079744 100644 --- a/python/sglang/srt/models/mllama4.py +++ b/python/sglang/srt/models/mllama4.py @@ -447,6 +447,8 @@ def forward( class Llama4ForConditionalGeneration(nn.Module): + verify_weights_on_load = True + packed_modules_mapping = { "qkv_proj": ["q_proj", "k_proj", "v_proj"], "gate_up_proj": ["gate_proj", "up_proj"], @@ -713,11 +715,7 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> Set[str]: loaded_params.add(name) self._handle_default_weight(name, loaded_weight, params_dict) - unloaded_params = params_dict.keys() - loaded_params - if unloaded_params: - logger.warning( - f"Some weights are not initialized from checkpoints {unloaded_params}" - ) + return loaded_params def _should_skip_weight(self, name: str) -> bool: """Check if we should skip loading this weight.""" 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..ebab43167a2c --- /dev/null +++ b/test/registered/unit/model_loader/test_load_and_verify_weights.py @@ -0,0 +1,87 @@ +# SPDX-License-Identifier: Apache-2.0 + +import unittest + +from sglang.srt.model_loader.weight_completeness import ( + load_and_verify_weights, + unloaded_required_params, +) +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +class _FakeModel: + def __init__( + self, + *, + loaded=None, + parameter_names=("required.weight", "optional.weight"), + verify_weights_on_load=True, + ): + self.loaded = loaded + self.parameter_names = parameter_names + self.verify_weights_on_load = verify_weights_on_load + self.seen_weights = None + + def load_weights(self, weights): + self.seen_weights = list(weights) + return self.loaded + + def named_parameters(self): + return ((name, object()) for name in self.parameter_names) + + @staticmethod + def is_optional_weight(name): + return name.startswith("optional.") + + +class TestUnloadedRequiredParams(unittest.TestCase): + def test_reports_only_unloaded_required_params(self): + self.assertEqual( + unloaded_required_params( + ["loaded.weight", "required.weight", "optional.weight"], + {"loaded.weight"}, + lambda name: name.startswith("optional."), + ), + {"required.weight"}, + ) + + +class TestLoadAndVerifyWeights(unittest.TestCase): + def test_full_load_verifies_migrated_model(self): + model = _FakeModel(loaded={"optional.weight"}) + + with self.assertRaisesRegex(RuntimeError, r"required\.weight"): + load_and_verify_weights(model, [("checkpoint.weight", object())]) + + self.assertEqual(model.seen_weights[0][0], "checkpoint.weight") + + def test_partial_load_can_skip_full_checkpoint_verification(self): + model = _FakeModel(loaded=set()) + + self.assertEqual( + load_and_verify_weights(model, [], is_full_checkpoint=False), + set(), + ) + + def test_unmigrated_model_preserves_none_return(self): + model = _FakeModel(loaded=None, verify_weights_on_load=False) + + self.assertIsNone(load_and_verify_weights(model, [])) + + def test_migrated_model_must_return_loaded_names(self): + model = _FakeModel(loaded=None) + + with self.assertRaisesRegex(TypeError, "must return loaded parameter names"): + load_and_verify_weights(model, []) + + def test_complete_load_returns_loaded_names(self): + loaded = {"required.weight"} + model = _FakeModel(loaded=loaded) + + self.assertIs(load_and_verify_weights(model, []), loaded) + + +if __name__ == "__main__": + unittest.main()