Skip to content

AutoWeightLoader support Sglang native models 1: demo - #28671

Merged
Fridge003 merged 8 commits into
sgl-project:mainfrom
JD-ETH:feat/auto-weight-loader-demo
Jul 21, 2026
Merged

Fridge003 merged 8 commits into
sgl-project:mainfrom
JD-ETH:feat/auto-weight-loader-demo

Conversation

@JD-ETH

@JD-ETH JD-ETH commented Jun 18, 2026 •

Copy link
Copy Markdown
Contributor

Motivation

AutoWeightLoader supports sglang models, better structural separation and guarantee.
Following discussion #24703 (comment)

Design and migration plan (PR1–PR9): #31051

Modifications

Speed Tests and Profiling

Checklist

  • local tests and ci test R.

CI States

Latest PR Test (Base): ❌ Run #29841391934
Latest PR Test (Extra): ❌ Run #29841391548

@JD-ETH
JD-ETH requested a review from b8zhong as a code owner June 18, 2026 18:59
@gemini-code-assist

Copy link
Copy Markdown
Contributor

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

@b8zhong b8zhong self-assigned this Jun 30, 2026
"""Decorator to register an architecture-specific weight remap function.

The decorated function receives a model instance and returns a WeightsMapper
(or None if no remap is needed for this configuration).

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we need it to return None? If no remap is needed, we'd better not use the decorator at all (to avoid confusion)

# ---------------------------------------------------------------------------


@dataclass

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we use msgspec.Struct instead (codebase rule)

@b8zhong b8zhong left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we show the use of WeightRemapRegistry in this PR (maybe modifying another model that requires this)

Comment thread python/sglang/srt/environ.py Outdated
Comment thread test/manual/test_weight_loader_v2_equiv.py Outdated
Comment thread test/manual/test_weight_loader_v2_equiv.py Outdated
Comment thread test/manual/test_weight_loader_v2_equiv.py
Comment thread test/manual/test_weight_loader_v2_equiv.py Outdated
Comment thread test/registered/unit/utils/test_auto_loader.py Outdated
Comment thread test/registered/unit/utils/test_auto_loader.py Outdated
Comment thread test/registered/unit/utils/test_auto_loader.py Outdated
Document PR1–PR9 coverage plan, post_load protocol, unsupported models,
and link from auto_loader module docstring.

Co-authored-by: Cursor <cursoragent@cursor.com>
@github-actions github-actions Bot added the documentation Improvements or additions to documentation label Jul 10, 2026
JD-ETH and others added 4 commits July 10, 2026 05:49
Replace P2P/streaming framing with verify_complete asserting all
expected runtime tensors were updated.

Co-authored-by: Cursor <cursoragent@cursor.com>
- Use msgspec.Struct for StackedParamsDispatch; RemapRegistry requires WeightsMapper
- Demo RemapRegistry on Llama v2 path with submodule stacked loaders
- Parametrize stacked/remap unit tests with real Llama checkpoint names
- Tighten v1/v2 equiv test: randomize init, compare all named_parameters
- Drop SGLANG_ENABLE_WEIGHT_LOADER_V2 env comment

Co-authored-by: Cursor <cursoragent@cursor.com>
Add registered SRTRunner tests for Qwen2 v1/v2 and transformers backend;
restore manual get_model state_dict equiv test; drop test_auto_loader.py.

Co-authored-by: Cursor <cursoragent@cursor.com>
Add wave-1 parallelism, sequencing for PR7–PR9, and infra conflict hotspots.

Co-authored-by: Cursor <cursoragent@cursor.com>
Comment thread test/manual/test_weight_loader_v2_equiv.py
Comment thread python/sglang/srt/model_loader/design.md Outdated
Remove design.md from the PR; point auto_loader docstring at the
migration tracker issue linked from RFC sgl-project#24703.

Co-authored-by: Cursor <cursoragent@cursor.com>
@Fridge003 Fridge003 added the run-ci CI: run the baseline test suite on this PR label Jul 16, 2026
@b8zhong

b8zhong commented Jul 21, 2026

Copy link
Copy Markdown
Collaborator

/rerun-test test/manual/test_weight_loader_v2_equiv.py test/registered/model_loading/test_weight_loader_v2_e2e.py

@github-actions

github-actions Bot commented Jul 21, 2026 •

Copy link
Copy Markdown
Contributor

Results for /rerun-test test/manual/test_weight_loader_v2_equiv.py test/registered/model_loading/test_weight_loader_v2_e2e.py:

🚀 1-gpu-5090 (1 test): ✅ View workflow run

cd test/ && python3 registered/model_loading/test_weight_loader_v2_e2e.py

⛔ test/manual/test_weight_loader_v2_equiv.py: No register_cuda_ci(runner_config=...) or register_cpu_ci() found in test/manual/test_weight_loader_v2_equiv.py. This file may not be a registered CI test.

@Fridge003
Fridge003 merged commit becf252 into sgl-project:main Jul 21, 2026
130 of 178 checks passed
Jialin added a commit to Jialin/sglang that referenced this pull request Jul 22, 2026
…gl-project#28671)

sgl-project#28671 refactored LlamaForCausalLM/Qwen2ForCausalLM.load_weights to dispatch to
self._legacy_load_weights / self._load_weights_v2. The reward and classification
wrappers (llama_reward, llama_classification, qwen2_rm, qwen2_classification)
borrow the base loader via an unbound call `Base.load_weights(self, ...)` but
inherit only nn.Module, so `self` has no `_legacy_load_weights` and weight load
fails with AttributeError (default path, since v2 is opt-in).

Call `_legacy_load_weights` directly, restoring the pre-sgl-project#28671 behavior. These
wrapped models are not AutoWeightLoader v2 citizens, so legacy is the correct path.
alexnails added a commit that referenced this pull request Jul 22, 2026
…#28671) (#31988)

Co-authored-by: Alex Nails <alex.nails@radixark.ai>
@JD-ETH

JD-ETH commented Jul 22, 2026

Copy link
Copy Markdown
Contributor Author

i will look into supporting the wrappers methods as well.

Zhylkaaa pushed a commit to Zhylkaaa/sglang that referenced this pull request Jul 29, 2026
Zhylkaaa pushed a commit to Zhylkaaa/sglang that referenced this pull request Jul 29, 2026
Atituiset pushed a commit to Atituiset/sglang that referenced this pull request Sep 10, 2026
Atituiset pushed a commit to Atituiset/sglang that referenced this pull request Sep 10, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

documentation Improvements or additions to documentation run-ci CI: run the baseline test suite on this PR

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants