AutoWeightLoader support Sglang native models 1: demo

python/sglang/srt/environ.py

@@ -207,6 +207,9 @@ class Envs:
     SGLANG_DISABLED_MODEL_ARCHS = EnvTuple(tuple())
     SGLANG_PREFETCH_BLOCK_SIZE_MB = EnvInt(16)
     SGLANG_GEMMA_OUT_OF_PLACE_POSITION_MUTATION = EnvBool(False)
+    # When enabled, migrated models use AutoWeightsLoader instead of
+    # their legacy hand-written load_weights loop.  Default off.
+    SGLANG_ENABLE_WEIGHT_LOADER_V2 = EnvBool(False)
 
     # Logging Options
     SGLANG_LOG_GC = EnvBool(False)

python/sglang/srt/model_loader/auto_loader.py

@@ -0,0 +1,208 @@
+# 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).
+- filter_pp_weights: generator that drops out-of-range PP layers.
+- WeightRemapRegistry: architecture-specific name remap registration.
+- Re-exports of AutoWeightsLoader and WeightsMapper from models/utils.py.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Iterable
+from dataclasses import dataclass, field
+from typing import Union
+
+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",
+    "filter_pp_weights",
+    "register_weight_remap",
+    "get_weight_remap",
+]
+
+
+# ---------------------------------------------------------------------------
+# Stacked Parameters Dispatch
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class StackedParamsDispatch:
+    """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)
+    """
+
+    # List of (fused_param_name, checkpoint_source_name, shard_id).
+    # shard_id is passed directly to param.weight_loader.
+    mappings: list[tuple[str, str, Union[int, str]]] = field(default_factory=list)
+
+    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),
+    ]
+)
+
+
+# ---------------------------------------------------------------------------
+# 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 | None]] = {}
+
+
+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
+    (or None if no remap is needed for this configuration).
+
+    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 | None]):
+        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/qwen2.py

@@ -105,6 +105,24 @@ def forward(
         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__(
@@ -195,6 +213,27 @@ def forward(
         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__(
@@ -567,6 +606,13 @@ def end_layer(self):
         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"),
@@ -636,6 +682,52 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
                 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