[Refactor] Move weight completeness verification into the loader

[EazyReal]

Upstream Review Target

The upstream PR for this branch is sgl-project#29323.

This fork-local PR mirrors the same branch and is kept only as a fork-side view of the final-state implementation.

Motivation

load_weights(weights) should be the reusable tensor-copy primitive. Full-checkpoint completeness is a loader concern because online/RL partial updates and pinned-cache reloads intentionally load only selected tensors.

Changes

• Add loader-side completeness verification for full checkpoint load paths.
• Keep verification opt-in with verify_weights_on_load.
• Let models declare legitimate full-checkpoint omissions with is_optional_weight(name).
• Migrate Gemma4, DeepSeek-V4 / DeepSeek-OCR / NextN, and Llama4 to return loaded names and opt in.
• Remove model-local missing-weight diagnostics that fire incorrectly on partial updates.

Verification

• git diff --check
• python3 -m py_compile on touched Python files
• focused ruff format --check
• focused ruff check on the new helper/test and changed DeepSeek NextN surface
• direct sanity check for unloaded_required_params

Summary by CodeRabbit

• New Features

  • Weight loading now includes automatic completeness checks, helping catch missing model weights earlier during startup.

• Bug Fixes

  • Improved handling of optional or expected-missing weights for several models, reducing false alarms during loading.
  • Streamlined loading behavior across multiple model formats so validation is applied more consistently.

• Tests

  • Added coverage for detecting missing required weights and ignoring optional ones.

---

CI States

Latest PR Test (Base): ❌ Run #30517416707
Latest PR Test (Extra): ❌ Run #30517416638

[coderabbitai]

📝 Walkthrough

Walkthrough

The loader now verifies loaded weights through a shared helper, and multiple load paths call it. Several model classes opt into verification or adjust optional-weight handling, while the new unit test covers missing and optional parameters.

Changes

Weight loading verification

Layer / File(s)	Summary
Helper and primary loader entrypoints python/sglang/srt/model_loader/loader.py, python/sglang/srt/model_loader/weight_completeness.py, test/registered/unit/model_loader/test_load_and_verify_weights.py	_load_and_verify_weights() centralizes load-and-verify behavior, BaseModelLoader and QuantizedRLModelLoader call it, and the new unit test covers missing, optional, and fully loaded parameters.
Alternate loader paths python/sglang/srt/model_loader/loader.py	BitsAndBytesModelLoader, GGUFModelLoader, RemoteModelLoader, and CPU quantization switch to the shared verification helper; nearby loader messages are reflowed in the same file.
Model opt-in and unload cleanup python/sglang/srt/models/deepseek_ocr.py, python/sglang/srt/models/deepseek_v4.py, python/sglang/srt/models/gemma4_causal.py, python/sglang/srt/models/gemma4_mm.py, python/sglang/srt/models/gemma4_unified.py	These models set verify_weights_on_load and remove local unloaded-weight validation and reporting; DeepseekV4ForCausalLM also adds is_optional_weight().
Optional weight handling python/sglang/srt/models/deepseek_v4_nextn.py, python/sglang/srt/models/mllama4.py	DeepseekV4ForCausalLMNextN returns the superclass load result and marks lm_head and embed_tokens as optional; Llama4ForConditionalGeneration sets verify_weights_on_load, remaps scale names through _remap_scale_name(), and skips absent remaps.

Sequence Diagram(s)

sequenceDiagram
  participant L as BaseModelLoader.load_weights_and_postprocess
  participant H as _load_and_verify_weights
  participant M as model.load_weights
  participant U as unloaded_required_params
  L->>H: pass model and weights
  H->>M: load_weights(weights)
  alt model.verify_weights_on_load
    H->>U: compute missing required params
    U-->>H: missing set
    alt missing set is not empty
      H-->>L: raise RuntimeError
    end
  end

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Possibly related issues

• [RFC] Weight loading: separate the load primitive from the completeness check #2: This PR adds verify_weights_on_load / is_optional_weight and centralizes completeness checks in the loader flow, matching the issue’s loader-side verification split.

Poem

I’m a rabbit with a tiny pen, 🐇
I hop through weights and count again.
Missing bits? I nudge them out;
Optional scales do not make me pout.
Hooray for loaders, hop, hop, hop!

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 16.22% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly and concisely describes the main change: moving weight completeness verification into the loader.
Description check	✅ Passed	The description explains the motivation, implementation changes, and verification steps, but omits the template checklist and review process sections.

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch vmax/weight-loading-loader-side

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@python/sglang/srt/model_loader/loader.py`:
- Around line 339-340: The completeness check in loader.py only verifies
model.named_parameters(), so it can miss required tensors that load_weights()
also populates from buffers. Update the verification path around
unloaded_required_params/model completeness to use the same tensor source
contract as load_weights(), including any buffer-backed entries merged into
params_dict by models like gemma4_mm and gemma4_causal. Then add a regression
test covering a required buffer-backed tensor to ensure it is reported as
missing when absent.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 0f304e31-50e5-4eb8-a19a-bc332a0d5fba

📥 Commits

Reviewing files that changed from the base of the PR and between 3344b73 and 36a4e9b68b1fdd96ce9ef6e6f544b6063fb8ec52.

📒 Files selected for processing (10)

• python/sglang/srt/model_loader/loader.py
• python/sglang/srt/model_loader/weight_completeness.py
• python/sglang/srt/models/deepseek_ocr.py
• python/sglang/srt/models/deepseek_v4.py
• python/sglang/srt/models/deepseek_v4_nextn.py
• python/sglang/srt/models/gemma4_causal.py
• python/sglang/srt/models/gemma4_mm.py
• python/sglang/srt/models/gemma4_unified.py
• python/sglang/srt/models/mllama4.py
• test/registered/unit/model_loader/test_load_and_verify_weights.py

💤 Files with no reviewable changes (1)

• python/sglang/srt/models/gemma4_unified.py