minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

AutoWeightLoader support Sglang native models 1: demo (#28671)

Browse Source
This commit is contained in:
JD
2026-07-21 14:32:13 -07:00
committed by GitHub
parent 0eae9423d8
commit becf252e6c
6 changed files with 592 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/environ.py
+1
View File
@@ -229,6 +229,7 @@ class Envs:
SGLANG_DISABLED_MODEL_ARCHS = EnvTuple(tuple())
SGLANG_PREFETCH_BLOCK_SIZE_MB = EnvInt(16)
SGLANG_GEMMA_OUT_OF_PLACE_POSITION_MUTATION = EnvBool(False)
SGLANG_ENABLE_WEIGHT_LOADER_V2 = EnvBool(False)
# HTTP server
# Decompress request bodies tagged with `x-body-compressed`.
python/sglang/srt/model_loader/auto_loader.py
+225
View File
@@ -0,0 +1,225 @@
# Copyright 2023-2025 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Centralized weight loading utilities for native SGLang models.
This module provides:
- StackedParamsDispatch: reusable stacked-parameter routing (qkv_proj, gate_up_proj).
- ExpertParamsDispatch: MoE expert_id + w1/w2/w3 shard routing.
- filter_pp_weights: generator that drops out-of-range PP layers.
- RemapRegistry: architecture-specific name remap registration.
- Re-exports of AutoWeightsLoader and WeightsMapper from models/utils.py.
Load / post-load split (PR1 protocol, see #31051):
load_weights(..., run_post_load=True) -> WeightLoadResult
post_load_weights(loaded=result) -> optional GPU derivations (MLA w_kc/w_vc, etc.)
Migration plan: https://github.com/sgl-project/sglang/issues/31051 (RFC #24703).
"""
from __future__ import annotations
from collections.abc import Callable, Iterable
from typing import Union
import msgspec
import torch
from torch import nn
from torch.nn import Parameter
from sglang.srt.layers.utils.common import get_layer_id
from sglang.srt.models.utils import AutoWeightsLoader, WeightsMapper
__all__ = [
"AutoWeightsLoader",
"WeightsMapper",
"StackedParamsDispatch",
"STANDARD_QKV_MAPPING",
"STANDARD_GATE_UP_MAPPING",
"STANDARD_STACKED_MAPPING",
"LLAMA_STACKED_MAPPING",
"filter_pp_weights",
"register_weight_remap",
"get_weight_remap",
]
# ---------------------------------------------------------------------------
# Stacked Parameters Dispatch
# ---------------------------------------------------------------------------
class StackedParamsDispatch(msgspec.Struct, frozen=True):
"""Centralized stacked-parameter loading for fused linear layers.
Handles the common pattern of mapping checkpoint names
(q_proj, k_proj, v_proj, gate_proj, up_proj) to fused runtime parameters
(qkv_proj, gate_up_proj) with the correct shard IDs.
Quantization is handled entirely by ``param.weight_loader`` on the layer —
this class only routes the tensor to the correct parameter with the correct
shard_id.
Usage::
mapping = StackedParamsDispatch([
("qkv_proj", "q_proj", "q"),
("qkv_proj", "k_proj", "k"),
("qkv_proj", "v_proj", "v"),
])
target = mapping.try_load(name, tensor, params_dict)
"""
# (fused_param_name, checkpoint_source_name, shard_id).
mappings: tuple[tuple[str, str, Union[int, str]], ...] = ()
def try_load(
self,
name: str,
tensor: torch.Tensor,
params_dict: dict[str, Parameter],
) -> str | None:
"""Try to load a weight via stacked mapping.
Returns the loaded runtime parameter name if matched and loaded,
the target name (for skip tracking) if the target param is missing
(e.g. optional bias), or None if no mapping matched.
"""
for fused_name, source_name, shard_id in self.mappings:
if source_name not in name:
continue
target = name.replace(source_name, fused_name)
param = params_dict.get(target)
if param is None:
# Parameter doesn't exist — e.g. GPTQ bias.
# Return target so caller can track the skip.
return target
param.weight_loader(param, tensor, shard_id)
return target
return None
# Pre-built instances for the most common decoder patterns.
STANDARD_QKV_MAPPING = StackedParamsDispatch(
mappings=(
("qkv_proj", "q_proj", "q"),
("qkv_proj", "k_proj", "k"),
("qkv_proj", "v_proj", "v"),
)
)
STANDARD_GATE_UP_MAPPING = StackedParamsDispatch(
mappings=(
("gate_up_proj", "gate_proj", 0),
("gate_up_proj", "up_proj", 1),
)
)
STANDARD_STACKED_MAPPING = StackedParamsDispatch(
mappings=(
("qkv_proj", "q_proj", "q"),
("qkv_proj", "k_proj", "k"),
("qkv_proj", "v_proj", "v"),
("gate_up_proj", "gate_proj", 0),
("gate_up_proj", "up_proj", 1),
)
)
# Llama-family full-path stacked mapping (dot-prefixed shard names).
LLAMA_STACKED_MAPPING = StackedParamsDispatch(
mappings=(
(".qkv_proj", ".q_proj", "q"),
(".qkv_proj", ".k_proj", "k"),
(".qkv_proj", ".v_proj", "v"),
(".gate_up_proj", ".gate_proj", 0),
(".gate_up_proj", ".up_proj", 1),
)
)
# ---------------------------------------------------------------------------
# Pipeline Parallel Weight Filter
# ---------------------------------------------------------------------------
def filter_pp_weights(
weights: Iterable[tuple[str, torch.Tensor]],
start_layer: int,
end_layer: int,
) -> Iterable[tuple[str, torch.Tensor]]:
"""Drop checkpoint entries whose layer index is outside [start_layer, end_layer).
Weights that don't contain a parseable layer index (embed_tokens, lm_head,
layer norms, etc.) are always passed through.
"""
for name, tensor in weights:
layer_id = get_layer_id(name)
if layer_id is not None and (layer_id < start_layer or layer_id >= end_layer):
continue
yield name, tensor
# ---------------------------------------------------------------------------
# Weight Remap Registry
# ---------------------------------------------------------------------------
_REMAP_REGISTRY: dict[str, Callable[[nn.Module], WeightsMapper]] = {}
def register_weight_remap(*class_names: str):
"""Decorator to register an architecture-specific weight remap function.
The decorated function receives a model instance and returns a
``WeightsMapper``. If no remap is needed for a model, do not register it.
Example::
@register_weight_remap("LlamaForCausalLM")
def _llama_remap(model) -> WeightsMapper:
return WeightsMapper(orig_to_new_suffix={
".activation_scale": ".input_scale",
".weight_scale_inv": ".weight_scale",
})
"""
def decorator(fn: Callable[[nn.Module], WeightsMapper]):
for cn in class_names:
_REMAP_REGISTRY[cn] = fn
return fn
return decorator
def get_weight_remap(model: nn.Module) -> WeightsMapper | None:
"""Get the registered weight remap for a model instance, or None."""
fn = _REMAP_REGISTRY.get(type(model).__name__)
if fn is None:
return None
return fn(model)
# ---------------------------------------------------------------------------
# Architecture-Specific Registrations
# ---------------------------------------------------------------------------
@register_weight_remap("LlamaForCausalLM")
def _llama_remap(model: nn.Module) -> WeightsMapper:
"""Llama-family FP8 scale suffix normalization."""
return WeightsMapper(
orig_to_new_suffix={
".activation_scale": ".input_scale",
".weight_scale_inv": ".weight_scale",
}
)
python/sglang/srt/models/llama.py
+77
View File
@@ -118,6 +118,22 @@ class LlamaMLP(nn.Module):
x, _ = self.down_proj(x)
return x
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> set[str]:
from sglang.srt.model_loader.auto_loader import STANDARD_GATE_UP_MAPPING
loaded: set[str] = set()
params_dict = dict(self.named_parameters())
for name, tensor in weights:
target = STANDARD_GATE_UP_MAPPING.try_load(name, tensor, params_dict)
if target is not None:
loaded.add(target)
continue
if name in params_dict:
wl = getattr(params_dict[name], "weight_loader", default_weight_loader)
wl(params_dict[name], tensor)
loaded.add(name)
return loaded
class LlamaAttention(nn.Module):
def __init__(
@@ -247,6 +263,22 @@ class LlamaAttention(nn.Module):
output, _ = self.o_proj(attn_output)
return output
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> set[str]:
from sglang.srt.model_loader.auto_loader import STANDARD_QKV_MAPPING
loaded: set[str] = set()
params_dict = dict(self.named_parameters())
for name, tensor in weights:
target = STANDARD_QKV_MAPPING.try_load(name, tensor, params_dict)
if target is not None:
loaded.add(target)
continue
if name in params_dict:
wl = getattr(params_dict[name], "weight_loader", default_weight_loader)
wl(params_dict[name], tensor)
loaded.add(name)
return loaded
class LlamaDecoderLayer(nn.Module):
def __init__(
@@ -623,6 +655,13 @@ class LlamaForCausalLM(nn.Module):
return len(params_dict)
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
from sglang.srt.environ import envs
if envs.SGLANG_ENABLE_WEIGHT_LOADER_V2.get():
return self._load_weights_v2(weights)
return self._legacy_load_weights(weights)
def _legacy_load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
stacked_params_mapping = [
# (param_name, shard_name, shard_id)
(".qkv_proj", ".q_proj", "q"),
@@ -695,6 +734,44 @@ class LlamaForCausalLM(nn.Module):
else:
logger.warning(f"Parameter {name} not found in params_dict")
def _load_weights_v2(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> set[str]:
"""AutoWeightsLoader path with RemapRegistry for FP8 suffix normalization."""
from sglang.srt.model_loader.auto_loader import (
AutoWeightsLoader,
filter_pp_weights,
get_weight_remap,
)
if hasattr(self.model, "start_layer"):
weights = filter_pp_weights(
weights, self.model.start_layer, self.model.end_layer
)
skip_prefixes = []
if self.config.tie_word_embeddings:
skip_prefixes.append("lm_head.")
loader = AutoWeightsLoader(
self,
skip_prefixes=skip_prefixes,
skip_substrs=["projector", "model.vision_tower"],
ignore_unexpected_suffixes=[".bias", ".kv_scale"],
)
mapper = get_weight_remap(self)
loaded = loader.load_weights(weights, mapper=mapper)
if self.config.tie_word_embeddings:
params_dict = dict(self.named_parameters())
if "lm_head.weight" in params_dict:
embed = dict(self.model.named_parameters()).get("embed_tokens.weight")
if embed is not None:
lm_head = params_dict["lm_head.weight"]
wl = getattr(lm_head, "weight_loader", default_weight_loader)
wl(lm_head, embed.data)
loaded.add("lm_head.weight")
return loaded
def get_weights_by_name(
self, name: str, truncate_size: int = 100, tp_size: int = 1
) -> Optional[torch.Tensor]:
python/sglang/srt/models/qwen2.py
+92
View File
@@ -104,6 +104,24 @@ class Qwen2MLP(nn.Module):
x, _ = self.down_proj(x, forward_batch=forward_batch)
return x
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> set[str]:
"""Load weights with centralized gate_up_proj stacked dispatch."""
from sglang.srt.model_loader.auto_loader import STANDARD_GATE_UP_MAPPING
loaded: set[str] = set()
params_dict = dict(self.named_parameters())
for name, tensor in weights:
target = STANDARD_GATE_UP_MAPPING.try_load(name, tensor, params_dict)
if target is not None:
loaded.add(target)
continue
# Direct params: down_proj, scales, etc.
if name in params_dict:
wl = getattr(params_dict[name], "weight_loader", default_weight_loader)
wl(params_dict[name], tensor)
loaded.add(name)
return loaded
class Qwen2Attention(nn.Module):
def __init__(
@@ -194,6 +212,27 @@ class Qwen2Attention(nn.Module):
output, _ = self.o_proj(attn_output)
return output
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> set[str]:
"""Load weights with centralized qkv_proj stacked dispatch."""
from sglang.srt.model_loader.auto_loader import STANDARD_QKV_MAPPING
loaded: set[str] = set()
params_dict = dict(self.named_parameters())
for name, tensor in weights:
target = STANDARD_QKV_MAPPING.try_load(name, tensor, params_dict)
if target is not None:
loaded.add(target)
continue
# Direct params: o_proj, scales, biases
if name in params_dict:
wl = getattr(params_dict[name], "weight_loader", default_weight_loader)
wl(params_dict[name], tensor)
loaded.add(name)
elif not (name.endswith(".bias") or name.endswith(".kv_scale")):
# Don't warn for optional bias or legacy kv_scale
pass
return loaded
class Qwen2DecoderLayer(nn.Module):
def __init__(
@@ -569,6 +608,13 @@ class Qwen2ForCausalLM(nn.Module):
return self.model.end_layer
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
from sglang.srt.environ import envs
if envs.SGLANG_ENABLE_WEIGHT_LOADER_V2.get():
return self._load_weights_v2(weights)
return self._legacy_load_weights(weights)
def _legacy_load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
stacked_params_mapping = [
# (param_name, shard_name, shard_id)
("qkv_proj", "q_proj", "q"),
@@ -638,6 +684,52 @@ class Qwen2ForCausalLM(nn.Module):
else:
logger.warning(f"Parameter {name} not found in params_dict")
def _load_weights_v2(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> set[str]:
"""AutoWeightsLoader-based weight loading.
Stacked params (qkv_proj, gate_up_proj) are handled by module-local
load_weights on Qwen2Attention and Qwen2MLP via walker delegation.
Quantization is handled entirely inside param.weight_loader.
"""
from sglang.srt.model_loader.auto_loader import (
AutoWeightsLoader,
filter_pp_weights,
)
# 1. PP layer filter — drop layers outside this rank's range.
if hasattr(self.model, "start_layer"):
weights = filter_pp_weights(
weights, self.model.start_layer, self.model.end_layer
)
# 2. Tied embeddings — skip lm_head from walker; we copy from
# embed_tokens after loading completes.
skip_prefixes = []
if self.config.tie_word_embeddings:
skip_prefixes.append("lm_head.")
# 3. Walk the module tree.
loader = AutoWeightsLoader(
self,
skip_prefixes=skip_prefixes,
skip_substrs=["projector", "model.vision_tower"],
ignore_unexpected_suffixes=[".bias", ".kv_scale"],
)
loaded = loader.load_weights(weights)
# 4. Tied embedding post-copy: replicate embed_tokens → lm_head.
if self.config.tie_word_embeddings:
params_dict = dict(self.named_parameters())
if "lm_head.weight" in params_dict:
embed = dict(self.model.named_parameters()).get("embed_tokens.weight")
if embed is not None:
lm_head = params_dict["lm_head.weight"]
wl = getattr(lm_head, "weight_loader", default_weight_loader)
wl(lm_head, embed.data)
loaded.add("lm_head.weight")
return loaded
def get_embed_and_head(self):
return self.model.embed_tokens.weight, self.lm_head.weight
test/manual/test_weight_loader_v2_equiv.py
+111
View File
@@ -0,0 +1,111 @@
# Copyright 2023-2025 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
# Manual verification for weight loader v2 (Qwen2 native path).
#
# Run:
# CUDA_VISIBLE_DEVICES=0 python test/manual/test_weight_loader_v2_equiv.py
#
# Engine-level e2e (Qwen2 + transformers backend) lives in:
# test/registered/model_loading/test_weight_loader_v2_e2e.py
import unittest
import torch
from sglang.srt.environ import envs
MODEL = "Qwen/Qwen2-0.5B"
def _init_model_parallel() -> None:
from sglang.srt.distributed import (
init_distributed_environment,
initialize_model_parallel,
)
from sglang.srt.distributed.parallel_state import monkey_patch_vllm_parallel_state
try:
init_distributed_environment(
backend="nccl",
world_size=1,
rank=0,
local_rank=0,
distributed_init_method="tcp://127.0.0.1:29634",
)
initialize_model_parallel(tensor_model_parallel_size=1)
monkey_patch_vllm_parallel_state()
except AssertionError:
pass
def _load_qwen2_native(v2: bool) -> torch.nn.Module:
from sglang.srt.configs.device_config import DeviceConfig
from sglang.srt.configs.load_config import LoadConfig
from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.model_loader import get_model
from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
from sglang.srt.utils import get_device
server_args = ServerArgs(
model_path=MODEL,
dtype=torch.float16,
trust_remote_code=True,
)
set_global_server_args_for_scheduler(server_args)
model_config = ModelConfig.from_server_args(server_args)
with envs.SGLANG_ENABLE_WEIGHT_LOADER_V2.override(v2):
return get_model(
model_config=model_config,
load_config=LoadConfig(),
device_config=DeviceConfig(get_device()),
)
def _state_dict_cpu(model: torch.nn.Module) -> dict[str, torch.Tensor]:
return {
name: param.detach().cpu().clone() for name, param in model.state_dict().items()
}
class TestWeightLoaderV2Equiv(unittest.TestCase):
@classmethod
def setUpClass(cls):
_init_model_parallel()
@unittest.skipIf(not torch.cuda.is_available(), "needs GPU")
def test_qwen2_v1_v2_state_dict_identical(self):
model_v1 = _load_qwen2_native(v2=False)
state_v1 = _state_dict_cpu(model_v1)
del model_v1
torch.cuda.empty_cache()
model_v2 = _load_qwen2_native(v2=True)
state_v2 = _state_dict_cpu(model_v2)
del model_v2
torch.cuda.empty_cache()
self.assertEqual(set(state_v1.keys()), set(state_v2.keys()))
for name in sorted(state_v1.keys()):
torch.testing.assert_close(
state_v1[name],
state_v2[name],
rtol=0,
atol=0,
msg=name,
)
if __name__ == "__main__":
unittest.main()
test/registered/model_loading/test_weight_loader_v2_e2e.py
+86
View File
@@ -0,0 +1,86 @@
# Copyright 2023-2025 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=120, stage="base-b", runner_config="1-gpu-small")
import multiprocessing as mp
import torch
from sglang.srt.environ import envs
from sglang.test.runners import SRTRunner, check_close_model_outputs
from sglang.test.test_utils import CustomTestCase
MODEL = "Qwen/Qwen2-0.5B"
SHORT_PROMPT = "The capital of the United Kingdom is"
class TestWeightLoaderV2E2E(CustomTestCase):
@classmethod
def setUpClass(cls):
mp.set_start_method("spawn", force=True)
def _runner_kwargs(self):
return dict(
torch_dtype=torch.float16,
model_type="generation",
disable_cuda_graph=True,
disable_radix_cache=True,
trust_remote_code=True,
max_total_tokens=2048,
)
def test_qwen2_native_v1_v2_generation_match(self):
prompts = [SHORT_PROMPT]
max_new_tokens = 32
kwargs = self._runner_kwargs()
with envs.SGLANG_ENABLE_WEIGHT_LOADER_V2.override(False):
with SRTRunner(MODEL, **kwargs) as runner_v1:
out_v1 = runner_v1.forward(prompts, max_new_tokens=max_new_tokens)
with envs.SGLANG_ENABLE_WEIGHT_LOADER_V2.override(True):
with SRTRunner(MODEL, **kwargs) as runner_v2:
out_v2 = runner_v2.forward(prompts, max_new_tokens=max_new_tokens)
check_close_model_outputs(
hf_outputs=out_v1,
srt_outputs=out_v2,
prefill_tolerance=1e-6,
decode_tolerance=1e-6,
rouge_l_tolerance=1.0,
debug_text="qwen2 native v1 vs v2 weight loader",
)
def test_transformers_impl_loads_and_generates(self):
prompts = [SHORT_PROMPT]
max_new_tokens = 16
with SRTRunner(
MODEL,
model_impl="transformers",
**self._runner_kwargs(),
) as runner:
outputs = runner.forward(prompts, max_new_tokens=max_new_tokens)
self.assertEqual(len(outputs.output_strs), 1)
self.assertGreater(len(outputs.output_strs[0]), 0)
if __name__ == "__main__":
import unittest
unittest.main()