diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 328164a94..3b03a8f15 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -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`. diff --git a/python/sglang/srt/model_loader/auto_loader.py b/python/sglang/srt/model_loader/auto_loader.py new file mode 100644 index 000000000..259c07dd4 --- /dev/null +++ b/python/sglang/srt/model_loader/auto_loader.py @@ -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", + } + ) diff --git a/python/sglang/srt/models/llama.py b/python/sglang/srt/models/llama.py index a7eec70af..6771c6136 100644 --- a/python/sglang/srt/models/llama.py +++ b/python/sglang/srt/models/llama.py @@ -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]: diff --git a/python/sglang/srt/models/qwen2.py b/python/sglang/srt/models/qwen2.py index 98ae54aeb..14e2601a7 100644 --- a/python/sglang/srt/models/qwen2.py +++ b/python/sglang/srt/models/qwen2.py @@ -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 diff --git a/test/manual/test_weight_loader_v2_equiv.py b/test/manual/test_weight_loader_v2_equiv.py new file mode 100644 index 000000000..1cf71c3e9 --- /dev/null +++ b/test/manual/test_weight_loader_v2_equiv.py @@ -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() diff --git a/test/registered/model_loading/test_weight_loader_v2_e2e.py b/test/registered/model_loading/test_weight_loader_v2_e2e.py new file mode 100644 index 000000000..7ec68c9c3 --- /dev/null +++ b/test/registered/model_loading/test_weight_loader_v2_e2e.py @@ -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()