From 1a3142d7e7974eb00a5abc1261540fa4b9890f54 Mon Sep 17 00:00:00 2001 From: JD-ETH Date: Thu, 18 Jun 2026 18:39:09 +0000 Subject: [PATCH 01/10] demo --- python/sglang/srt/environ.py | 3 + python/sglang/srt/model_loader/auto_loader.py | 208 ++++++++ python/sglang/srt/models/qwen2.py | 92 ++++ test/manual/test_weight_loader_v2_equiv.py | 118 +++++ .../registered/unit/utils/test_auto_loader.py | 483 ++++++++++++++++++ 5 files changed, 904 insertions(+) create mode 100644 python/sglang/srt/model_loader/auto_loader.py create mode 100644 test/manual/test_weight_loader_v2_equiv.py create mode 100644 test/registered/unit/utils/test_auto_loader.py diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 1694e651c773..cc24806220b4 100644 --- a/python/sglang/srt/environ.py +++ b/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) 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 000000000000..518a64e97304 --- /dev/null +++ b/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", + } + ) diff --git a/python/sglang/srt/models/qwen2.py b/python/sglang/srt/models/qwen2.py index b7ddc653601d..301ccc210652 100644 --- a/python/sglang/srt/models/qwen2.py +++ b/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 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 000000000000..4afe048d16b7 --- /dev/null +++ b/test/manual/test_weight_loader_v2_equiv.py @@ -0,0 +1,118 @@ +# 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. +# ============================================================================== +"""Differential validation: v1 vs v2 weight loaders produce identical state. + +This test loads the same weights through both the legacy and AutoWeightsLoader +paths, then asserts bit-identical parameter state. + +Run with: + SGLANG_IS_IN_CI=1 python -m pytest test/manual/test_weight_loader_v2_equiv.py -v + +Requires: GPU, small model downloadable from HF. +""" + +import unittest + +import torch + +from sglang.srt.environ import envs + + +class TestWeightLoaderV2Equivalence(unittest.TestCase): + """Prove v1 and v2 weight loaders produce bit-identical model state.""" + + MODEL_PATH = "Qwen/Qwen2-0.5B" + + @unittest.skipIf(not torch.cuda.is_available(), "needs GPU") + def test_qwen2_v1_v2_equivalence(self): + from transformers import AutoConfig + + from sglang.srt.model_loader.weight_utils import get_model_weights_iter + + # Load weights once into memory + config = AutoConfig.from_pretrained(self.MODEL_PATH, trust_remote_code=True) + + # We need to instantiate two Qwen2 models and load them separately. + # For this test we use a simplified approach: load weights list once, + # then replay through both paths. + from sglang.srt.models.qwen2 import Qwen2ForCausalLM + + # Minimal construction args (TP=1, no quant, no PP) + with torch.device("cuda"): + model_v1 = Qwen2ForCausalLM(config) + model_v2 = Qwen2ForCausalLM(config) + + # Collect all weights + weights = list(get_model_weights_iter(self.MODEL_PATH)) + + # Load with v1 (legacy) + with envs.SGLANG_ENABLE_WEIGHT_LOADER_V2.override(False): + model_v1.load_weights(iter(weights)) + + # Load with v2 (AutoWeightsLoader) + with envs.SGLANG_ENABLE_WEIGHT_LOADER_V2.override(True): + loaded_names = model_v2.load_weights(iter(weights)) + + # loaded_names should be non-empty + self.assertIsNotNone(loaded_names) + self.assertGreater(len(loaded_names), 0) + + # Compare all parameters + v1_params = dict(model_v1.named_parameters()) + v2_params = dict(model_v2.named_parameters()) + + self.assertEqual( + set(v1_params.keys()), + set(v2_params.keys()), + "Parameter name sets differ between v1 and v2", + ) + + mismatches = [] + for name in sorted(v1_params.keys()): + p1 = v1_params[name] + p2 = v2_params[name] + if not torch.equal(p1.data, p2.data): + max_diff = (p1.data - p2.data).abs().max().item() + mismatches.append(f" {name}: max_diff={max_diff:.6e}") + + if mismatches: + self.fail( + f"Parameter mismatches ({len(mismatches)}):\n" + + "\n".join(mismatches[:20]) + ) + + # Compare all buffers + v1_buffers = dict(model_v1.named_buffers()) + v2_buffers = dict(model_v2.named_buffers()) + + for name in v1_buffers: + if name not in v2_buffers: + continue + b1 = v1_buffers[name] + b2 = v2_buffers[name] + if not torch.equal(b1, b2): + max_diff = (b1 - b2).abs().max().item() + mismatches.append(f" buffer {name}: max_diff={max_diff:.6e}") + + if mismatches: + self.fail( + f"Buffer mismatches ({len(mismatches)}):\n" + "\n".join(mismatches[:20]) + ) + + print(f"✓ V1 and V2 produce identical state for {self.MODEL_PATH}") + print(f"✓ V2 loaded {len(loaded_names)} parameter names") + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/utils/test_auto_loader.py b/test/registered/unit/utils/test_auto_loader.py new file mode 100644 index 000000000000..6376dbc5e57b --- /dev/null +++ b/test/registered/unit/utils/test_auto_loader.py @@ -0,0 +1,483 @@ +# 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. +# ============================================================================== +"""Unit tests for the centralized weight loading utilities (auto_loader.py). + +These tests are CPU-only and don't require model weights or GPU. +""" + +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + +import unittest +from unittest.mock import MagicMock + +import torch +from torch import nn +from torch.nn import Parameter + +from sglang.srt.model_loader.auto_loader import ( + STANDARD_GATE_UP_MAPPING, + STANDARD_QKV_MAPPING, + STANDARD_STACKED_MAPPING, + filter_pp_weights, + get_weight_remap, + register_weight_remap, +) +from sglang.srt.models.utils import AutoWeightsLoader, WeightsMapper + +# --------------------------------------------------------------------------- +# StackedParamsDispatch Tests +# --------------------------------------------------------------------------- + + +class TestStackedParamsDispatch(unittest.TestCase): + def _make_param_with_loader(self, shape): + """Create a parameter with a mock weight_loader.""" + param = Parameter(torch.zeros(shape)) + param.weight_loader = MagicMock() + return param + + def test_qkv_mapping_matches_q_proj(self): + params = {"qkv_proj.weight": self._make_param_with_loader((30, 10))} + tensor = torch.ones(10, 10) + + result = STANDARD_QKV_MAPPING.try_load("q_proj.weight", tensor, params) + + self.assertEqual(result, "qkv_proj.weight") + params["qkv_proj.weight"].weight_loader.assert_called_once_with( + params["qkv_proj.weight"], tensor, "q" + ) + + def test_qkv_mapping_matches_k_proj(self): + params = {"qkv_proj.weight": self._make_param_with_loader((30, 10))} + tensor = torch.ones(10, 10) + + result = STANDARD_QKV_MAPPING.try_load("k_proj.weight", tensor, params) + + self.assertEqual(result, "qkv_proj.weight") + params["qkv_proj.weight"].weight_loader.assert_called_once_with( + params["qkv_proj.weight"], tensor, "k" + ) + + def test_qkv_mapping_matches_v_proj(self): + params = {"qkv_proj.weight": self._make_param_with_loader((30, 10))} + tensor = torch.ones(10, 10) + + result = STANDARD_QKV_MAPPING.try_load("v_proj.weight", tensor, params) + + self.assertEqual(result, "qkv_proj.weight") + params["qkv_proj.weight"].weight_loader.assert_called_once_with( + params["qkv_proj.weight"], tensor, "v" + ) + + def test_gate_up_mapping_matches_gate_proj(self): + params = {"gate_up_proj.weight": self._make_param_with_loader((20, 10))} + tensor = torch.ones(10, 10) + + result = STANDARD_GATE_UP_MAPPING.try_load("gate_proj.weight", tensor, params) + + self.assertEqual(result, "gate_up_proj.weight") + params["gate_up_proj.weight"].weight_loader.assert_called_once_with( + params["gate_up_proj.weight"], tensor, 0 + ) + + def test_gate_up_mapping_matches_up_proj(self): + params = {"gate_up_proj.weight": self._make_param_with_loader((20, 10))} + tensor = torch.ones(10, 10) + + result = STANDARD_GATE_UP_MAPPING.try_load("up_proj.weight", tensor, params) + + self.assertEqual(result, "gate_up_proj.weight") + params["gate_up_proj.weight"].weight_loader.assert_called_once_with( + params["gate_up_proj.weight"], tensor, 1 + ) + + def test_no_match_returns_none(self): + params = {"qkv_proj.weight": self._make_param_with_loader((30, 10))} + tensor = torch.ones(10, 10) + + result = STANDARD_QKV_MAPPING.try_load("o_proj.weight", tensor, params) + + self.assertIsNone(result) + params["qkv_proj.weight"].weight_loader.assert_not_called() + + def test_missing_target_param_returns_target_name(self): + # Parameter doesn't exist in params_dict (e.g. GPTQ bias) + params = {} # No qkv_proj.bias + tensor = torch.ones(10) + + result = STANDARD_QKV_MAPPING.try_load("q_proj.bias", tensor, params) + + # Returns target name for skip tracking + self.assertEqual(result, "qkv_proj.bias") + + def test_standard_stacked_mapping_covers_all(self): + params = { + "qkv_proj.weight": self._make_param_with_loader((30, 10)), + "gate_up_proj.weight": self._make_param_with_loader((20, 10)), + } + tensor = torch.ones(10, 10) + + # All 5 source names should match + for source, expected_target in [ + ("q_proj.weight", "qkv_proj.weight"), + ("k_proj.weight", "qkv_proj.weight"), + ("v_proj.weight", "qkv_proj.weight"), + ("gate_proj.weight", "gate_up_proj.weight"), + ("up_proj.weight", "gate_up_proj.weight"), + ]: + result = STANDARD_STACKED_MAPPING.try_load(source, tensor, params) + self.assertEqual(result, expected_target, f"Failed for {source}") + + +# --------------------------------------------------------------------------- +# WeightsMapper Tests +# --------------------------------------------------------------------------- + + +class TestWeightsMapper(unittest.TestCase): + def test_substr_remap(self): + mapper = WeightsMapper(orig_to_new_substr={".old_name.": ".new_name."}) + weights = [("prefix.old_name.weight", torch.ones(2, 2))] + + result = list(mapper.apply(weights)) + + self.assertEqual(len(result), 1) + self.assertEqual(result[0][0], "prefix.new_name.weight") + + def test_prefix_remap(self): + mapper = WeightsMapper(orig_to_new_prefix={"model.model.": "model."}) + weights = [("model.model.layers.0.weight", torch.ones(2, 2))] + + result = list(mapper.apply(weights)) + + self.assertEqual(len(result), 1) + self.assertEqual(result[0][0], "model.layers.0.weight") + + def test_suffix_remap(self): + mapper = WeightsMapper(orig_to_new_suffix={".activation_scale": ".input_scale"}) + weights = [("layer.activation_scale", torch.ones(2))] + + result = list(mapper.apply(weights)) + + self.assertEqual(len(result), 1) + self.assertEqual(result[0][0], "layer.input_scale") + + def test_drop_via_none(self): + mapper = WeightsMapper(orig_to_new_substr={"rotary_emb": None}) + weights = [ + ("layer.rotary_emb.inv_freq", torch.ones(2)), + ("layer.o_proj.weight", torch.ones(2, 2)), + ] + + result = list(mapper.apply(weights)) + + self.assertEqual(len(result), 1) + self.assertEqual(result[0][0], "layer.o_proj.weight") + + def test_composition_or(self): + mapper1 = WeightsMapper(orig_to_new_substr={".old1.": ".new1."}) + mapper2 = WeightsMapper(orig_to_new_suffix={".scale_inv": ".scale"}) + + composed = mapper1 | mapper2 + weights = [ + ("prefix.old1.weight", torch.ones(2, 2)), + ("layer.scale_inv", torch.ones(2)), + ] + + result = list(composed.apply(weights)) + + self.assertEqual(result[0][0], "prefix.new1.weight") + self.assertEqual(result[1][0], "layer.scale") + + def test_longest_substr_match_wins(self): + mapper = WeightsMapper( + orig_to_new_substr={ + "proj": "proj_replaced", + "q_proj": "qkv_proj", # Longer match should win + } + ) + weights = [("layer.q_proj.weight", torch.ones(2, 2))] + + result = list(mapper.apply(weights)) + + self.assertEqual(result[0][0], "layer.qkv_proj.weight") + + +# --------------------------------------------------------------------------- +# AutoWeightsLoader Tests +# --------------------------------------------------------------------------- + + +class TestAutoWeightsLoader(unittest.TestCase): + def test_direct_param_load(self): + model = nn.Linear(4, 4, bias=False) + loader = AutoWeightsLoader(model) + tensor = torch.ones(4, 4) + + loaded = loader.load_weights([("weight", tensor)]) + + self.assertIn("weight", loaded) + self.assertTrue(torch.equal(model.weight.data, tensor)) + + def test_nested_module_load(self): + class Model(nn.Module): + def __init__(self): + super().__init__() + self.layer = nn.Linear(4, 4, bias=False) + + model = Model() + loader = AutoWeightsLoader(model) + tensor = torch.ones(4, 4) + + loaded = loader.load_weights([("layer.weight", tensor)]) + + self.assertIn("layer.weight", loaded) + self.assertTrue(torch.equal(model.layer.weight.data, tensor)) + + def test_skip_prefixes(self): + class Model(nn.Module): + def __init__(self): + super().__init__() + self.kept = nn.Linear(4, 4, bias=False) + self.skipped = nn.Linear(4, 4, bias=False) + + model = Model() + loader = AutoWeightsLoader(model, skip_prefixes=["skipped."]) + tensor = torch.ones(4, 4) + + loaded = loader.load_weights( + [ + ("kept.weight", tensor), + ("skipped.weight", tensor * 2), + ] + ) + + self.assertIn("kept.weight", loaded) + self.assertNotIn("skipped.weight", loaded) + self.assertTrue(torch.equal(model.kept.weight.data, tensor)) + # skipped should remain at init value (zeros-like from default) + + def test_skip_rotary_embeddings(self): + class Model(nn.Module): + def __init__(self): + super().__init__() + self.weight = nn.Parameter(torch.zeros(4)) + + model = Model() + loader = AutoWeightsLoader(model) + + loaded = loader.load_weights( + [ + ("weight", torch.ones(4)), + ("rotary_emb.inv_freq", torch.ones(10)), + ] + ) + + self.assertIn("weight", loaded) + self.assertNotIn("rotary_emb.inv_freq", loaded) + + def test_module_delegation(self): + """Walker delegates to child module's load_weights if present.""" + + class ChildModule(nn.Module): + def __init__(self): + super().__init__() + self.param = nn.Parameter(torch.zeros(4)) + + def load_weights(self, weights): + loaded = set() + for name, tensor in weights: + if name == "source_name": + self.param.data.copy_(tensor) + loaded.add("param") + return loaded + + class Model(nn.Module): + def __init__(self): + super().__init__() + self.child = ChildModule() + + model = Model() + loader = AutoWeightsLoader(model) + tensor = torch.ones(4) + + loaded = loader.load_weights([("child.source_name", tensor)]) + + self.assertIn("child.param", loaded) + self.assertTrue(torch.equal(model.child.param.data, tensor)) + + def test_ignore_unexpected_suffix(self): + class Model(nn.Module): + def __init__(self): + super().__init__() + self.layer = nn.Linear(4, 4, bias=False) + + model = Model() + loader = AutoWeightsLoader(model, ignore_unexpected_suffixes=[".bias"]) + + # Should not raise for unexpected .bias weight nested under layer + loaded = loader.load_weights( + [ + ("layer.weight", torch.ones(4, 4)), + ("layer.bias", torch.ones(4)), # model's layer has no bias + ] + ) + + self.assertIn("layer.weight", loaded) + + def test_raises_on_truly_unexpected(self): + model = nn.Linear(4, 4, bias=False) + loader = AutoWeightsLoader(model) + + with self.assertRaises(ValueError): + loader.load_weights([("nonexistent_param", torch.ones(4))]) + + def test_mapper_applied(self): + model = nn.Linear(4, 4, bias=False) + mapper = WeightsMapper(orig_to_new_substr={"old_weight": "weight"}) + loader = AutoWeightsLoader(model) + tensor = torch.ones(4, 4) + + loaded = loader.load_weights([("old_weight", tensor)], mapper=mapper) + + self.assertIn("weight", loaded) + self.assertTrue(torch.equal(model.weight.data, tensor)) + + def test_pp_missing_layer_skipped(self): + """PPMissingLayer modules should be silently skipped.""" + + class PPMissingLayer(nn.Module): + pass + + class Model(nn.Module): + def __init__(self): + super().__init__() + self.layer = PPMissingLayer() + self.embed = nn.Linear(4, 4, bias=False) + + model = Model() + loader = AutoWeightsLoader(model) + + # Weights for the missing layer should not raise + loaded = loader.load_weights( + [ + ("embed.weight", torch.ones(4, 4)), + ("layer.weight", torch.ones(4, 4)), + ] + ) + + self.assertIn("embed.weight", loaded) + + +# --------------------------------------------------------------------------- +# filter_pp_weights Tests +# --------------------------------------------------------------------------- + + +class TestFilterPPWeights(unittest.TestCase): + def test_drops_out_of_range_layers(self): + weights = [ + ("model.layers.0.self_attn.weight", torch.ones(4)), + ("model.layers.1.self_attn.weight", torch.ones(4)), + ("model.layers.2.self_attn.weight", torch.ones(4)), + ("model.layers.3.self_attn.weight", torch.ones(4)), + ] + + result = list(filter_pp_weights(weights, start_layer=1, end_layer=3)) + + self.assertEqual(len(result), 2) + self.assertEqual(result[0][0], "model.layers.1.self_attn.weight") + self.assertEqual(result[1][0], "model.layers.2.self_attn.weight") + + def test_passes_non_layer_weights(self): + weights = [ + ("model.embed_tokens.weight", torch.ones(4)), + ("lm_head.weight", torch.ones(4)), + ("model.layers.0.self_attn.weight", torch.ones(4)), + ("model.norm.weight", torch.ones(4)), + ] + + result = list(filter_pp_weights(weights, start_layer=1, end_layer=3)) + + # embed_tokens, lm_head, norm pass through (no layer id) + # layer 0 is filtered out + names = [name for name, _ in result] + self.assertIn("model.embed_tokens.weight", names) + self.assertIn("lm_head.weight", names) + self.assertIn("model.norm.weight", names) + self.assertNotIn("model.layers.0.self_attn.weight", names) + + def test_empty_range_drops_all_layers(self): + weights = [ + ("model.layers.0.weight", torch.ones(4)), + ("model.embed_tokens.weight", torch.ones(4)), + ] + + result = list(filter_pp_weights(weights, start_layer=5, end_layer=5)) + + names = [name for name, _ in result] + self.assertNotIn("model.layers.0.weight", names) + self.assertIn("model.embed_tokens.weight", names) + + +# --------------------------------------------------------------------------- +# WeightRemapRegistry Tests +# --------------------------------------------------------------------------- + + +class TestWeightRemapRegistry(unittest.TestCase): + def test_registered_model_returns_mapper(self): + # LlamaForCausalLM is registered in auto_loader.py + class FakeLlama(nn.Module): + pass + + FakeLlama.__name__ = "LlamaForCausalLM" + model = FakeLlama() + + mapper = get_weight_remap(model) + + self.assertIsNotNone(mapper) + self.assertIsInstance(mapper, WeightsMapper) + + def test_unknown_model_returns_none(self): + class UnknownModel(nn.Module): + pass + + model = UnknownModel() + + mapper = get_weight_remap(model) + + self.assertIsNone(mapper) + + def test_custom_registration(self): + @register_weight_remap("TestModelXYZ") + def _test_remap(model): + return WeightsMapper(orig_to_new_substr={"foo": "bar"}) + + class TestModelXYZ(nn.Module): + pass + + model = TestModelXYZ() + mapper = get_weight_remap(model) + + self.assertIsNotNone(mapper) + weights = [("prefix.foo.weight", torch.ones(2))] + result = list(mapper.apply(weights)) + self.assertEqual(result[0][0], "prefix.bar.weight") + + +if __name__ == "__main__": + unittest.main() From 2ba329218d1eee9a4fa05d053a3f3d666472384f Mon Sep 17 00:00:00 2001 From: JD-ETH Date: Fri, 10 Jul 2026 05:43:16 +0000 Subject: [PATCH 02/10] Add weight loader v2 migration design doc. MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Document PR1–PR9 coverage plan, post_load protocol, unsupported models, and link from auto_loader module docstring. Co-authored-by: Cursor --- python/sglang/srt/model_loader/auto_loader.py | 9 +- python/sglang/srt/model_loader/design.md | 326 ++++++++++++++++++ 2 files changed, 334 insertions(+), 1 deletion(-) create mode 100644 python/sglang/srt/model_loader/design.md diff --git a/python/sglang/srt/model_loader/auto_loader.py b/python/sglang/srt/model_loader/auto_loader.py index 518a64e97304..68e363f08158 100644 --- a/python/sglang/srt/model_loader/auto_loader.py +++ b/python/sglang/srt/model_loader/auto_loader.py @@ -15,9 +15,16 @@ 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. -- WeightRemapRegistry: architecture-specific name remap registration. +- RemapRegistry: architecture-specific name remap registration. - Re-exports of AutoWeightsLoader and WeightsMapper from models/utils.py. + +Load / post-load split (PR1 protocol, see design.md): + load_weights(..., run_post_load=True) -> WeightLoadResult + post_load_weights(loaded=result) -> optional GPU derivations (MLA w_kc/w_vc, etc.) + +Migration plan: design.md in this directory. """ from __future__ import annotations diff --git a/python/sglang/srt/model_loader/design.md b/python/sglang/srt/model_loader/design.md new file mode 100644 index 000000000000..aa7283c13c36 --- /dev/null +++ b/python/sglang/srt/model_loader/design.md @@ -0,0 +1,326 @@ +# Weight Loader v2 — design and migration plan + +Tracking issue: [RFC #24703](https://github.com/sgl-project/sglang/issues/24703) +Demo PR: [#28671](https://github.com/sgl-project/sglang/pull/28671) + +--- + +## Why this design + +**Complete loading, no stragglers.** Every checkpoint tensor that enters the loader and is not +explicitly skipped must produce a `WeightLoadRecord`, or `verify_complete` raises +`IncompleteWeightLoadError`. That replaces silent `logger.warning` drops and ad-hoc expert-shard +checks scattered across individual model files. + +**Clean, visible weight mapping.** Each load is recorded as +`(checkpoint_name → runtime_name, shard_id, expert_id)`. Stacked qkv/gate_up and MoE expert +shards are first-class — not buried inside 70-line per-model loops. Mapping tables live in +`StackedParamsDispatch`, `ExpertParamsDispatch`, and `RemapRegistry` instead of being +copy-pasted 116 times across `srt/models/`. + +**Enables advanced weight flows (e.g. P2P / streaming updates).** A mergeable `WeightLoadResult` +with per-shard records lets callers assert that a partial update batch finished, accumulate +across multiple `load_weights` calls, and know exactly which runtime slot each HF tensor +touched. The load / `post_load_weights` split separates tensor placement from optional +GPU-specific derivation (MLA `w_kc`/`w_vc`, etc.), which advanced transports need to control +independently. + +All migrations stay behind `SGLANG_ENABLE_WEIGHT_LOADER_V2` (opt-in, default off). Legacy +`load_weights` paths remain until explicitly removed per model. + +--- + +## Migration overview (~180 `load_weights` files) + +Cumulative auto-loader coverage by PR (unique model implementation files, ~180 total): + +| PR | Adds (files) | Cumulative | Coverage | +|----|-------------:|-----------:|---------:| +| **PR1** | 2 (`qwen2.py`, `transformers.py`) | 2 | 1% | +| **PR2** | ~40 | ~42 | 23% | +| **PR3** | ~22 | ~64 | 36% | +| **PR4** | ~5 | ~69 | 38% | +| **PR5** | 1 | ~70 | 39% | +| **PR6** | 1 | ~71 | 39% | +| **PR7** | 1 | ~72 | 40% | +| **PR8** | ~25 | ~97 | 54% | +| **PR9** | ~35 wrappers | ~132 | 73% | + +`transformers.py` uses `AutoWeightsLoader` today (no v2 gate); counted in PR1 as already on the +walker stack. PR9 wrapper files add no new loader logic — they inherit v2 from a migrated +parent — but are included in the cumulative file count once smoke-tested. + +### Protocol extensions (introduced across the PR stack) + +| Extension | PR | Purpose | +|-----------|-----|---------| +| `post_load_weights` split | PR1 / PR4 | Separate tensor load from GPU derivations | +| Shared-expert preload hook | PR4 | Rename `shared_experts` → routed expert slot before dispatch | +| MLA `RemapRegistry` entries | PR4 | `fused_qkv_a_proj_with_mqa`, scale remaps | +| `FusedExpertDispatch` | PR5 | One fused ckpt key → multiple expert shard records | +| Packed-shard fan-out | PR5 | One ckpt tensor → multiple column shard records | +| Hybrid mixer submodule loaders | PR6 | Mamba + MoE in one decoder block | +| Custom loader decomposition | PR7 | Large bespoke loops → submodule overrides | +| Vision / audio tower `load_weights` | PR8 | Walker delegates to per-tower loaders | + +--- + +## PR1 — Infrastructure + first model (this PR) + +**Addresses:** Core abstractions and the reference dense migration pattern. + +**Coverage:** +2 files (`qwen2.py`, `transformers.py` on walker stack) (~1%). + +**Work:** + +| Area | Deliverable | +|------|-------------| +| `auto_loader.py` | `WeightLoadRecord`, `WeightLoadResult`, `StackedParamsDispatch`, `ExpertParamsDispatch`, `RemapRegistry`, `verify_complete` | +| `models/utils.py` | `AutoWeightsLoader` returns `WeightLoadResult` | +| Load protocol | `load_weights(..., run_post_load=True)` + default no-op `post_load_weights()` | +| `models/qwen2.py` | v2 gate; `Qwen2Attention` / `Qwen2MLP` submodule loaders | +| `models/transformers.py` | already on `AutoWeightsLoader`; document as covered (no v2 gate) | +| `environ.py` | `SGLANG_ENABLE_WEIGHT_LOADER_V2 = EnvBool(False)` | +| Tests | `test_auto_loader.py`, `test_auto_loader_moe.py`, `test_weight_loader_v2_equiv.py` | + +Quant `process_weights_after_loading` remains in `loader.py` (unchanged). + +--- + +## PR2 — Standard dense models (~41 files) + +**Addresses:** Models with `stacked_params_mapping` only (qkv + gate_up), no expert mapping, +no MLA mixin. Mechanical migration: env gate, submodule loaders, optional `RemapRegistry`. + +**Coverage:** +~40 files (~23% cumulative). + +**Work per file:** Rename loop → `_legacy_load_weights`; add `_load_weights_v2` using +`AutoWeightsLoader` + `StackedParamsDispatch` on attention/MLP submodules; register FP8/kv +remaps where needed. + +**Models:** + +| File | Notes | +|------|-------| +| `qwen2.py` | done in PR1 | +| `qwen3.py` | + `maybe_remap_kv_scale` | +| `qwen3_classification.py` | inherits Qwen3 | +| `llama.py` | FP8 suffix + kv scale via `RemapRegistry` | +| `mistral.py` | delegates → Llama | +| `glm4.py` | dot-prefix stacked | +| `gemma.py`, `gemma2.py`, `gemma3_causal.py`, `gemma3n_causal.py` | gemma2/3 + kv remap | +| `olmo.py`, `olmo2.py`, `internlm2.py`, `granite.py` | standard | +| `stablelm.py`, `starcoder2.py`, `opt.py`, `orion.py` | standard | +| `baichuan.py`, `xverse.py`, `exaone.py`, `exaone4.py` | standard / remap | +| `commandr.py`, `apertus.py`, `arcee.py`, `nemotron_nas.py` | + remap | +| `falcon_h1.py`, `lfm2.py`, `iquest_loopcoder.py`, `jet_nemotron.py` | standard | +| `mimo.py`, `gpt_j.py`, `sdar.py`, `dflash.py`, `afmoe.py` | standard / remap | +| `bert.py`, `roberta.py`, `clip.py`, `whisper.py`, `siglip2.py` | encoders | +| `llama_embedding.py`, `torch_native_llama.py` | variants | + +Cherry-pick walker diffs from +[brayden/weight-loader-refactor-part-1](https://github.com/sgl-project/sglang/compare/main...brayden/weight-loader-refactor-part-1) +where possible. + +**Tests:** parametrized CPU mapping tests; spot-check GPU equiv (`Llama-3.2-1B`, `Qwen3-1.7B-FP8`). + +--- + +## PR3 — Standard MoE models (~22 files) + +**Addresses:** Models using `FusedMoE.make_expert_params_mapping` with dense stacked layers. + +**Coverage:** +~22 files (~36% cumulative). + +**Work per file:** `ExpertParamsDispatch` on `FusedMoE` submodule; dense submodules as PR2; +`verify_expert_shards_complete` where all routed expert shards are mandatory (e.g. `laguna.py`). + +**Models:** + +| File | Notes | +|------|-------| +| `qwen2_moe.py`, `qwen3_moe.py` | standard MoE | +| `mixtral.py`, `glm4_moe.py`, `olmoe.py`, `phimoe.py` | standard | +| `cohere2_moe.py`, `exaone_moe.py`, `lfm2_moe.py` | standard | +| `hunyuan.py`, `hunyuan_v3.py`, `bailing_moe.py` | standard | +| `laguna.py` | + expert-shard completeness assert | +| `ernie4.py`, `grok.py`, `kimi_linear.py`, `llada2.py` | standard | +| `minimax_m2.py`, `sdar_moe.py`, `sarvam_moe.py`, `dbrx.py` | standard | +| `xverse_moe.py`, `granitemoe.py` | standard / delegates Mixtral | + +**Tests:** `Qwen/Qwen1.5-MoE-A2.7B` equiv; extend `test_auto_loader_moe.py`. + +--- + +## PR4 — MLA mixin models (~5 files) + +**Addresses:** `DeepseekV2WeightLoaderMixin` / `do_load_weights` — shared expert preload, MLA +name remaps, and `post_load_weights` for `w_kc`/`w_vc` derivation. + +**Coverage:** +~5 files (~38% cumulative). + +**Work:** Replace mixin loop with walker + submodule loaders; wire shared-expert preload hook; +split `post_load_weights` from tensor load (PR1 protocol); shrink `deepseek_weight_loader.py`. + +**Models:** + +| File | Notes | +|------|-------| +| `deepseek_v2.py` | V2, V3, V32 | +| `deepseek_common/deepseek_weight_loader.py` | mixin refactor | +| `glm4_moe_lite.py` | GLM-4.7 | +| `bailing_moe_linear.py` | linear MoE variant | +| `kimi_k25_eagle3.py` | eagle MLA draft | + +**Tests:** `deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct`; GLM-4.7 mapping units. + +--- + +## PR5 — `qwen3_5.py` + +**Addresses:** Fused MoE checkpoint format (`gate_up_proj` → w1 + w3 shards) and packed linear +shard fan-out. + +**Coverage:** +1 file (~39% cumulative). + +**Work:** Add `FusedExpertDispatch`; fan-out multi-shard records from single ckpt tensors; +migrate `Qwen3_5MoeForConditionalGeneration`, `Qwen3_5ForConditionalGeneration`. + +**Tests:** `Qwen/Qwen3.5-4B`; fused-expert unit tests. + +--- + +## PR6 — `qwen3_next.py` + +**Addresses:** Hybrid mamba + MoE decoder — mixer and expert paths need separate submodule +`load_weights`. + +**Coverage:** +1 file (~39% cumulative). + +**Work:** Submodule overrides on mamba mixer and MoE block; standard dispatch elsewhere. + +**Tests:** mapping unit tests (full-model GPU equiv impractical at default sizes). + +--- + +## PR7 — `deepseek_v4.py` + +**Addresses:** Large bespoke loader (~335 lines) with custom expert routing and `post_load_weights`. + +**Coverage:** +1 file (~40% cumulative). + +**Work:** Decompose into submodule loaders; reuse PR4 post_load protocol. + +**Tests:** unit mapping tests; targeted manual validation. + +--- + +## PR8 — Multimodal models (~25 files) + +**Addresses:** Vision, audio, and OCR tower weights routed alongside text trunk. + +**Coverage:** +~25 files (~54% cumulative). + +**Work:** Text trunk uses PR2/PR3 walker; add `VisionTower.load_weights` / +`AudioTower.load_weights` submodule delegation; `skip_substrs` for optional towers. Reuse +`FusedExpertDispatch` from PR5 where VL-MoE uses fused expert ckpts. + +**Models (representative):** `qwen2_vl.py`, `qwen2_5_vl.py`, `qwen3_vl.py`, `qwen3_vl_moe.py`, +`qwen3_omni_moe.py`, `qwen2_audio.py`, `qwen3_asr.py`, `deepseek_vl2.py`, `deepseek_ocr.py`, +`deepseek_janus_pro.py`, `glm4v.py`, `glm4v_moe.py`, `glm_ocr.py`, `glmasr.py`, `minicpmv.py`, +`internvl.py`, `kimi_vl.py`, `moss_vl.py`, `gemma3_mm.py`, `gemma4_mm.py`, `pixtral.py`, … + +**Tests:** `Qwen/Qwen2.5-VL-3B-Instruct`, `Qwen/Qwen3-VL-2B-Thinking`; per-tower CPU units. + +--- + +## PR9 — Wrapper models (~35 files) + +**Addresses:** Thin classes that only call `super().load_weights` (eagle, MTP, NextN, +classification, reward, delegates). + +**Coverage:** +~35 wrapper files (~73% cumulative). No new loader implementation — inherits +v2 from parent. Smoke-test that `super().load_weights` forwards `WeightLoadResult` and +`run_post_load` correctly; fix NextN `__init__` bypass mapping gaps. + +**Work:** Ensure mappings are available when `__init__` is bypassed (NextN pattern); forward +`run_post_load` to parent; smoke tests only. + +**Models (representative):** `qwen2_eagle.py`, `qwen3_moe_mtp.py`, `qwen3_5_mtp.py`, +`qwen3_next_mtp.py`, `deepseek_nextn.py`, `deepseek_v4_nextn.py`, `llava.py`, `granitemoe.py`, +`mistral.py`, `qwen2_rm.py`, … + +--- + +## Not supported today (open to community contribution) + +The PR1–PR9 plan does **not** include these models. They stay on legacy `load_weights` until +someone contributes a dedicated migration. v2 remains opt-in via `SGLANG_ENABLE_WEIGHT_LOADER_V2`. + +### Legacy architectures + +| File | Reason | +|------|--------| +| `qwen.py` | Superseded by Qwen2+; hand-rolled loader | +| `deepseek.py` | Superseded by DeepSeek V2+; hand-rolled loader | + +### Custom loaders (no `stacked_params_mapping` pattern) + +| File | Reason | +|------|--------| +| `phi.py` | `packed_modules_mapping` driven | +| `phi3_small.py` | Custom loader loop | +| `solar.py` | `packed_modules_mapping` driven | +| `gpt2.py` | Direct param iteration | +| `gpt_bigcode.py` | Direct param iteration | +| `chatglm.py` | Custom GLM QKV routing | +| `minicpm.py` | Manual w1/w2/w3 expert table | +| `minicpm3.py` | Manual expert mapping | +| `parakeet.py` | Custom loader | +| `persimmon.py` | Custom loader | + +### Bespoke loaders (need dedicated design per model) + +| File | Reason | +|------|--------| +| `step3p5.py` | Custom expert name matching (~146L loader) | +| `step3_vl.py` | Custom expert matching + vision | +| `step3p5_mtp.py` | MTP variant of step3 loader | +| `step3_vl_10b.py` | VL delegate + step3 patterns | +| `step3p7.py` | Step family bespoke routing | +| `nemotron_h.py` | Non-standard expert `weight_loader` signature | +| `nemotron_h_mtp.py` | MTP wrapper on nemotron_h | +| `zaya.py` | MOD routing / custom MoE | +| `longcat_flash.py` | Hybrid loader (~180L) | +| `longcat_flash_nextn.py` | NextN + longcat patterns | +| `mimo_v2.py` | Custom qkv routing (~198L loader) | +| `mimo_v2_nextn.py` | MTP on mimo_v2 | +| `mimo_v2_asr.py` | ASR delegate on mimo_v2 | +| `gpt_oss.py` | Unique quant routing | +| `gemma4_causal.py` | Fused expert ckpt (needs `FusedExpertDispatch` beyond PR5 scope) | +| `granitemoehybrid.py` | Mamba + MoE hybrid (~150L loader) | +| `mindspore.py` | MindSpore cell loader path | + +Pull requests that migrate any of the above to the v2 stack (walker + dispatch + records) are +welcome; add mapping unit tests and a manual or CI equiv test where a small checkpoint exists. + +--- + +## Checklist + +- [ ] PR1: infra + post_load protocol + `qwen2.py` + `transformers.py` (walker) +- [ ] PR2: standard dense (~41 files) +- [ ] PR3: standard MoE (~22 files) +- [ ] PR4: MLA mixin (~5 files) +- [ ] PR5: `qwen3_5.py` +- [ ] PR6: `qwen3_next.py` +- [ ] PR7: `deepseek_v4.py` +- [ ] PR8: multimodal (~25 files) +- [ ] PR9: wrappers (~35 files) + +--- + +## References + +- [brayden/weight-loader-refactor-part-1](https://github.com/sgl-project/sglang/compare/main...brayden/weight-loader-refactor-part-1) — cherry-pick dense walker diffs into PR2 +- [PR #24256](https://github.com/sgl-project/sglang/pull/24256) — checksum for production equiv validation From 33a3b4719a9bb80dd2cfe9add2d30d78067cc6ab Mon Sep 17 00:00:00 2001 From: JD-ETH Date: Fri, 10 Jul 2026 05:49:15 +0000 Subject: [PATCH 03/10] Reframe design benefits around native completeness check. Replace P2P/streaming framing with verify_complete asserting all expected runtime tensors were updated. Co-authored-by: Cursor --- python/sglang/srt/model_loader/design.md | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/python/sglang/srt/model_loader/design.md b/python/sglang/srt/model_loader/design.md index aa7283c13c36..d5f316081bba 100644 --- a/python/sglang/srt/model_loader/design.md +++ b/python/sglang/srt/model_loader/design.md @@ -18,12 +18,11 @@ shards are first-class — not buried inside 70-line per-model loops. Mapping ta `StackedParamsDispatch`, `ExpertParamsDispatch`, and `RemapRegistry` instead of being copy-pasted 116 times across `srt/models/`. -**Enables advanced weight flows (e.g. P2P / streaming updates).** A mergeable `WeightLoadResult` -with per-shard records lets callers assert that a partial update batch finished, accumulate -across multiple `load_weights` calls, and know exactly which runtime slot each HF tensor -touched. The load / `post_load_weights` split separates tensor placement from optional -GPU-specific derivation (MLA `w_kc`/`w_vc`, etc.), which advanced transports need to control -independently. +**Native weight-loading completeness check.** v2 loads call `verify_complete` to assert every +expected runtime tensor was updated — no silent skips, no stragglers. A mergeable +`WeightLoadResult` with per-shard records supports this check across single or batched +`load_weights` calls. The load / `post_load_weights` split applies completeness verification +after placement and before optional GPU-specific derivation (MLA `w_kc`/`w_vc`, etc.). All migrations stay behind `SGLANG_ENABLE_WEIGHT_LOADER_V2` (opt-in, default off). Legacy `load_weights` paths remain until explicitly removed per model. From c0bcbec63503ff140d49e78f425870fae0b2fe6c Mon Sep 17 00:00:00 2001 From: JD-ETH Date: Fri, 10 Jul 2026 18:08:30 +0000 Subject: [PATCH 04/10] Address PR review comments on weight loader v2 demo. - 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 --- python/sglang/srt/environ.py | 2 - python/sglang/srt/model_loader/auto_loader.py | 42 ++- python/sglang/srt/models/llama.py | 77 +++++ test/manual/test_weight_loader_v2_equiv.py | 52 +-- .../registered/unit/utils/test_auto_loader.py | 299 ++++++------------ 5 files changed, 214 insertions(+), 258 deletions(-) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index cc24806220b4..e0d68bdfeb9d 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -207,8 +207,6 @@ 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 diff --git a/python/sglang/srt/model_loader/auto_loader.py b/python/sglang/srt/model_loader/auto_loader.py index 68e363f08158..7bcb8e2eb2a6 100644 --- a/python/sglang/srt/model_loader/auto_loader.py +++ b/python/sglang/srt/model_loader/auto_loader.py @@ -30,9 +30,9 @@ from __future__ import annotations from collections.abc import Callable, Iterable -from dataclasses import dataclass, field from typing import Union +import msgspec import torch from torch import nn from torch.nn import Parameter @@ -47,6 +47,7 @@ "STANDARD_QKV_MAPPING", "STANDARD_GATE_UP_MAPPING", "STANDARD_STACKED_MAPPING", + "LLAMA_STACKED_MAPPING", "filter_pp_weights", "register_weight_remap", "get_weight_remap", @@ -58,8 +59,7 @@ # --------------------------------------------------------------------------- -@dataclass -class StackedParamsDispatch: +class StackedParamsDispatch(msgspec.Struct, frozen=True): """Centralized stacked-parameter loading for fused linear layers. Handles the common pattern of mapping checkpoint names @@ -80,9 +80,8 @@ class StackedParamsDispatch: 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) + # (fused_param_name, checkpoint_source_name, shard_id). + mappings: tuple[tuple[str, str, Union[int, str]], ...] = () def try_load( self, @@ -113,28 +112,39 @@ def try_load( # Pre-built instances for the most common decoder patterns. STANDARD_QKV_MAPPING = StackedParamsDispatch( - mappings=[ + mappings=( ("qkv_proj", "q_proj", "q"), ("qkv_proj", "k_proj", "k"), ("qkv_proj", "v_proj", "v"), - ] + ) ) STANDARD_GATE_UP_MAPPING = StackedParamsDispatch( - mappings=[ + mappings=( ("gate_up_proj", "gate_proj", 0), ("gate_up_proj", "up_proj", 1), - ] + ) ) STANDARD_STACKED_MAPPING = StackedParamsDispatch( - mappings=[ + 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), + ) ) @@ -164,14 +174,14 @@ def filter_pp_weights( # Weight Remap Registry # --------------------------------------------------------------------------- -_REMAP_REGISTRY: dict[str, Callable[[nn.Module], WeightsMapper | None]] = {} +_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 - (or None if no remap is needed for this configuration). + The decorated function receives a model instance and returns a + ``WeightsMapper``. If no remap is needed for a model, do not register it. Example:: @@ -183,7 +193,7 @@ def _llama_remap(model) -> WeightsMapper: }) """ - def decorator(fn: Callable[[nn.Module], WeightsMapper | None]): + def decorator(fn: Callable[[nn.Module], WeightsMapper]): for cn in class_names: _REMAP_REGISTRY[cn] = fn return fn diff --git a/python/sglang/srt/models/llama.py b/python/sglang/srt/models/llama.py index 742f810accd6..d62309813694 100644 --- a/python/sglang/srt/models/llama.py +++ b/python/sglang/srt/models/llama.py @@ -123,6 +123,22 @@ def forward( ) 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__( @@ -252,6 +268,22 @@ def forward( 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__( @@ -628,6 +660,13 @@ def get_num_params(self): 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"), @@ -700,6 +739,44 @@ 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 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/test/manual/test_weight_loader_v2_equiv.py b/test/manual/test_weight_loader_v2_equiv.py index 4afe048d16b7..cd8aa0856a3a 100644 --- a/test/manual/test_weight_loader_v2_equiv.py +++ b/test/manual/test_weight_loader_v2_equiv.py @@ -11,16 +11,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Differential validation: v1 vs v2 weight loaders produce identical state. - -This test loads the same weights through both the legacy and AutoWeightsLoader -paths, then asserts bit-identical parameter state. - -Run with: - SGLANG_IS_IN_CI=1 python -m pytest test/manual/test_weight_loader_v2_equiv.py -v - -Requires: GPU, small model downloadable from HF. -""" +"""Differential validation: v1 vs v2 weight loaders produce identical state.""" import unittest @@ -30,8 +21,6 @@ class TestWeightLoaderV2Equivalence(unittest.TestCase): - """Prove v1 and v2 weight loaders produce bit-identical model state.""" - MODEL_PATH = "Qwen/Qwen2-0.5B" @unittest.skipIf(not torch.cuda.is_available(), "needs GPU") @@ -39,48 +28,31 @@ def test_qwen2_v1_v2_equivalence(self): from transformers import AutoConfig from sglang.srt.model_loader.weight_utils import get_model_weights_iter + from sglang.srt.models.qwen2 import Qwen2ForCausalLM - # Load weights once into memory config = AutoConfig.from_pretrained(self.MODEL_PATH, trust_remote_code=True) - # We need to instantiate two Qwen2 models and load them separately. - # For this test we use a simplified approach: load weights list once, - # then replay through both paths. - from sglang.srt.models.qwen2 import Qwen2ForCausalLM - - # Minimal construction args (TP=1, no quant, no PP) with torch.device("cuda"): model_v1 = Qwen2ForCausalLM(config) model_v2 = Qwen2ForCausalLM(config) - # Collect all weights + torch.manual_seed(0) + for param in (*model_v1.parameters(), *model_v2.parameters()): + param.data.random_() + weights = list(get_model_weights_iter(self.MODEL_PATH)) - # Load with v1 (legacy) with envs.SGLANG_ENABLE_WEIGHT_LOADER_V2.override(False): model_v1.load_weights(iter(weights)) - # Load with v2 (AutoWeightsLoader) with envs.SGLANG_ENABLE_WEIGHT_LOADER_V2.override(True): loaded_names = model_v2.load_weights(iter(weights)) - # loaded_names should be non-empty - self.assertIsNotNone(loaded_names) self.assertGreater(len(loaded_names), 0) - # Compare all parameters - v1_params = dict(model_v1.named_parameters()) v2_params = dict(model_v2.named_parameters()) - - self.assertEqual( - set(v1_params.keys()), - set(v2_params.keys()), - "Parameter name sets differ between v1 and v2", - ) - mismatches = [] - for name in sorted(v1_params.keys()): - p1 = v1_params[name] + for name, p1 in model_v1.named_parameters(): p2 = v2_params[name] if not torch.equal(p1.data, p2.data): max_diff = (p1.data - p2.data).abs().max().item() @@ -92,14 +64,13 @@ def test_qwen2_v1_v2_equivalence(self): + "\n".join(mismatches[:20]) ) - # Compare all buffers + self.assertTrue(loaded_names.issubset(set(v2_params.keys()))) + v1_buffers = dict(model_v1.named_buffers()) v2_buffers = dict(model_v2.named_buffers()) - - for name in v1_buffers: + for name, b1 in v1_buffers.items(): if name not in v2_buffers: continue - b1 = v1_buffers[name] b2 = v2_buffers[name] if not torch.equal(b1, b2): max_diff = (b1 - b2).abs().max().item() @@ -110,9 +81,6 @@ def test_qwen2_v1_v2_equivalence(self): f"Buffer mismatches ({len(mismatches)}):\n" + "\n".join(mismatches[:20]) ) - print(f"✓ V1 and V2 produce identical state for {self.MODEL_PATH}") - print(f"✓ V2 loaded {len(loaded_names)} parameter names") - if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/utils/test_auto_loader.py b/test/registered/unit/utils/test_auto_loader.py index 6376dbc5e57b..b1fa978b7cca 100644 --- a/test/registered/unit/utils/test_auto_loader.py +++ b/test/registered/unit/utils/test_auto_loader.py @@ -11,10 +11,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Unit tests for the centralized weight loading utilities (auto_loader.py). - -These tests are CPU-only and don't require model weights or GPU. -""" +"""Unit tests for the centralized weight loading utilities (auto_loader.py).""" from sglang.test.ci.ci_register import register_cpu_ci @@ -23,6 +20,7 @@ import unittest from unittest.mock import MagicMock +import pytest import torch from torch import nn from torch.nn import Parameter @@ -37,188 +35,113 @@ ) from sglang.srt.models.utils import AutoWeightsLoader, WeightsMapper -# --------------------------------------------------------------------------- -# StackedParamsDispatch Tests -# --------------------------------------------------------------------------- - - -class TestStackedParamsDispatch(unittest.TestCase): - def _make_param_with_loader(self, shape): - """Create a parameter with a mock weight_loader.""" - param = Parameter(torch.zeros(shape)) - param.weight_loader = MagicMock() - return param - - def test_qkv_mapping_matches_q_proj(self): - params = {"qkv_proj.weight": self._make_param_with_loader((30, 10))} - tensor = torch.ones(10, 10) - - result = STANDARD_QKV_MAPPING.try_load("q_proj.weight", tensor, params) - - self.assertEqual(result, "qkv_proj.weight") - params["qkv_proj.weight"].weight_loader.assert_called_once_with( - params["qkv_proj.weight"], tensor, "q" - ) - - def test_qkv_mapping_matches_k_proj(self): - params = {"qkv_proj.weight": self._make_param_with_loader((30, 10))} - tensor = torch.ones(10, 10) - - result = STANDARD_QKV_MAPPING.try_load("k_proj.weight", tensor, params) - - self.assertEqual(result, "qkv_proj.weight") - params["qkv_proj.weight"].weight_loader.assert_called_once_with( - params["qkv_proj.weight"], tensor, "k" - ) - - def test_qkv_mapping_matches_v_proj(self): - params = {"qkv_proj.weight": self._make_param_with_loader((30, 10))} - tensor = torch.ones(10, 10) - - result = STANDARD_QKV_MAPPING.try_load("v_proj.weight", tensor, params) - - self.assertEqual(result, "qkv_proj.weight") - params["qkv_proj.weight"].weight_loader.assert_called_once_with( - params["qkv_proj.weight"], tensor, "v" - ) - - def test_gate_up_mapping_matches_gate_proj(self): - params = {"gate_up_proj.weight": self._make_param_with_loader((20, 10))} - tensor = torch.ones(10, 10) - - result = STANDARD_GATE_UP_MAPPING.try_load("gate_proj.weight", tensor, params) - - self.assertEqual(result, "gate_up_proj.weight") - params["gate_up_proj.weight"].weight_loader.assert_called_once_with( - params["gate_up_proj.weight"], tensor, 0 - ) - - def test_gate_up_mapping_matches_up_proj(self): - params = {"gate_up_proj.weight": self._make_param_with_loader((20, 10))} - tensor = torch.ones(10, 10) - - result = STANDARD_GATE_UP_MAPPING.try_load("up_proj.weight", tensor, params) - - self.assertEqual(result, "gate_up_proj.weight") - params["gate_up_proj.weight"].weight_loader.assert_called_once_with( - params["gate_up_proj.weight"], tensor, 1 - ) - - def test_no_match_returns_none(self): - params = {"qkv_proj.weight": self._make_param_with_loader((30, 10))} - tensor = torch.ones(10, 10) - - result = STANDARD_QKV_MAPPING.try_load("o_proj.weight", tensor, params) - - self.assertIsNone(result) - params["qkv_proj.weight"].weight_loader.assert_not_called() - - def test_missing_target_param_returns_target_name(self): - # Parameter doesn't exist in params_dict (e.g. GPTQ bias) - params = {} # No qkv_proj.bias - tensor = torch.ones(10) - - result = STANDARD_QKV_MAPPING.try_load("q_proj.bias", tensor, params) - - # Returns target name for skip tracking - self.assertEqual(result, "qkv_proj.bias") - - def test_standard_stacked_mapping_covers_all(self): - params = { - "qkv_proj.weight": self._make_param_with_loader((30, 10)), - "gate_up_proj.weight": self._make_param_with_loader((20, 10)), - } - tensor = torch.ones(10, 10) - - # All 5 source names should match - for source, expected_target in [ - ("q_proj.weight", "qkv_proj.weight"), - ("k_proj.weight", "qkv_proj.weight"), - ("v_proj.weight", "qkv_proj.weight"), - ("gate_proj.weight", "gate_up_proj.weight"), - ("up_proj.weight", "gate_up_proj.weight"), - ]: - result = STANDARD_STACKED_MAPPING.try_load(source, tensor, params) - self.assertEqual(result, expected_target, f"Failed for {source}") - -# --------------------------------------------------------------------------- -# WeightsMapper Tests -# --------------------------------------------------------------------------- +def _make_param_with_loader(shape): + param = Parameter(torch.zeros(shape)) + param.weight_loader = MagicMock() + return param -class TestWeightsMapper(unittest.TestCase): - def test_substr_remap(self): - mapper = WeightsMapper(orig_to_new_substr={".old_name.": ".new_name."}) - weights = [("prefix.old_name.weight", torch.ones(2, 2))] - - result = list(mapper.apply(weights)) - - self.assertEqual(len(result), 1) - self.assertEqual(result[0][0], "prefix.new_name.weight") - - def test_prefix_remap(self): - mapper = WeightsMapper(orig_to_new_prefix={"model.model.": "model."}) - weights = [("model.model.layers.0.weight", torch.ones(2, 2))] - - result = list(mapper.apply(weights)) - - self.assertEqual(len(result), 1) - self.assertEqual(result[0][0], "model.layers.0.weight") - - def test_suffix_remap(self): - mapper = WeightsMapper(orig_to_new_suffix={".activation_scale": ".input_scale"}) - weights = [("layer.activation_scale", torch.ones(2))] - - result = list(mapper.apply(weights)) - - self.assertEqual(len(result), 1) - self.assertEqual(result[0][0], "layer.input_scale") - - def test_drop_via_none(self): - mapper = WeightsMapper(orig_to_new_substr={"rotary_emb": None}) - weights = [ - ("layer.rotary_emb.inv_freq", torch.ones(2)), - ("layer.o_proj.weight", torch.ones(2, 2)), - ] - - result = list(mapper.apply(weights)) - - self.assertEqual(len(result), 1) - self.assertEqual(result[0][0], "layer.o_proj.weight") - - def test_composition_or(self): - mapper1 = WeightsMapper(orig_to_new_substr={".old1.": ".new1."}) - mapper2 = WeightsMapper(orig_to_new_suffix={".scale_inv": ".scale"}) - - composed = mapper1 | mapper2 - weights = [ - ("prefix.old1.weight", torch.ones(2, 2)), - ("layer.scale_inv", torch.ones(2)), - ] - - result = list(composed.apply(weights)) +@pytest.mark.parametrize( + "mapping,source_name,expected_target,expected_shard", + [ + (STANDARD_QKV_MAPPING, "q_proj.weight", "qkv_proj.weight", "q"), + (STANDARD_QKV_MAPPING, "k_proj.weight", "qkv_proj.weight", "k"), + (STANDARD_QKV_MAPPING, "v_proj.weight", "qkv_proj.weight", "v"), + (STANDARD_GATE_UP_MAPPING, "gate_proj.weight", "gate_up_proj.weight", 0), + (STANDARD_GATE_UP_MAPPING, "up_proj.weight", "gate_up_proj.weight", 1), + ], +) +def test_stacked_params_dispatch(mapping, source_name, expected_target, expected_shard): + params = {expected_target: _make_param_with_loader((30, 10))} + tensor = torch.ones(10, 10) + + result = mapping.try_load(source_name, tensor, params) + + assert result == expected_target + params[expected_target].weight_loader.assert_called_once_with( + params[expected_target], tensor, expected_shard + ) + + +def test_stacked_params_dispatch_no_match_returns_none(): + params = {"qkv_proj.weight": _make_param_with_loader((30, 10))} + tensor = torch.ones(10, 10) + + result = STANDARD_QKV_MAPPING.try_load("o_proj.weight", tensor, params) + + assert result is None + params["qkv_proj.weight"].weight_loader.assert_not_called() + + +def test_stacked_params_dispatch_missing_target_param_returns_target_name(): + params = {} + tensor = torch.ones(10) + + result = STANDARD_QKV_MAPPING.try_load("q_proj.bias", tensor, params) + + assert result == "qkv_proj.bias" + + +def test_standard_stacked_mapping_covers_all(): + params = { + "qkv_proj.weight": _make_param_with_loader((30, 10)), + "gate_up_proj.weight": _make_param_with_loader((20, 10)), + } + tensor = torch.ones(10, 10) + + for source, expected_target in [ + ("q_proj.weight", "qkv_proj.weight"), + ("k_proj.weight", "qkv_proj.weight"), + ("v_proj.weight", "qkv_proj.weight"), + ("gate_proj.weight", "gate_up_proj.weight"), + ("up_proj.weight", "gate_up_proj.weight"), + ]: + result = STANDARD_STACKED_MAPPING.try_load(source, tensor, params) + assert result == expected_target, f"Failed for {source}" + + +@pytest.mark.parametrize( + "checkpoint_name,expected_name", + [ + ( + "model.layers.0.self_attn.q_proj.activation_scale", + "model.layers.0.self_attn.q_proj.input_scale", + ), + ( + "model.layers.0.mlp.down_proj.weight_scale_inv", + "model.layers.0.mlp.down_proj.weight_scale", + ), + ( + "model.layers.3.self_attn.k_proj.weight_scale_inv", + "model.layers.3.self_attn.k_proj.weight_scale", + ), + ], +) +def test_llama_remap_registry(checkpoint_name, expected_name): + class FakeLlama(nn.Module): + pass - self.assertEqual(result[0][0], "prefix.new1.weight") - self.assertEqual(result[1][0], "layer.scale") + FakeLlama.__name__ = "LlamaForCausalLM" + mapper = get_weight_remap(FakeLlama()) + assert mapper is not None - def test_longest_substr_match_wins(self): - mapper = WeightsMapper( - orig_to_new_substr={ - "proj": "proj_replaced", - "q_proj": "qkv_proj", # Longer match should win - } - ) - weights = [("layer.q_proj.weight", torch.ones(2, 2))] + result = list(mapper.apply([(checkpoint_name, torch.ones(2))])) + assert len(result) == 1 + assert result[0][0] == expected_name - result = list(mapper.apply(weights)) - self.assertEqual(result[0][0], "layer.qkv_proj.weight") +def test_weights_mapper_drops_rotary_emb(): + mapper = WeightsMapper(orig_to_new_substr={"rotary_emb": None}) + weights = [ + ("layer.rotary_emb.inv_freq", torch.ones(2)), + ("model.layers.0.self_attn.o_proj.weight", torch.ones(2, 2)), + ] + result = list(mapper.apply(weights)) -# --------------------------------------------------------------------------- -# AutoWeightsLoader Tests -# --------------------------------------------------------------------------- + assert len(result) == 1 + assert result[0][0] == "model.layers.0.self_attn.o_proj.weight" class TestAutoWeightsLoader(unittest.TestCase): @@ -268,7 +191,6 @@ def __init__(self): self.assertIn("kept.weight", loaded) self.assertNotIn("skipped.weight", loaded) self.assertTrue(torch.equal(model.kept.weight.data, tensor)) - # skipped should remain at init value (zeros-like from default) def test_skip_rotary_embeddings(self): class Model(nn.Module): @@ -290,8 +212,6 @@ def __init__(self): self.assertNotIn("rotary_emb.inv_freq", loaded) def test_module_delegation(self): - """Walker delegates to child module's load_weights if present.""" - class ChildModule(nn.Module): def __init__(self): super().__init__() @@ -328,11 +248,10 @@ def __init__(self): model = Model() loader = AutoWeightsLoader(model, ignore_unexpected_suffixes=[".bias"]) - # Should not raise for unexpected .bias weight nested under layer loaded = loader.load_weights( [ ("layer.weight", torch.ones(4, 4)), - ("layer.bias", torch.ones(4)), # model's layer has no bias + ("layer.bias", torch.ones(4)), ] ) @@ -357,8 +276,6 @@ def test_mapper_applied(self): self.assertTrue(torch.equal(model.weight.data, tensor)) def test_pp_missing_layer_skipped(self): - """PPMissingLayer modules should be silently skipped.""" - class PPMissingLayer(nn.Module): pass @@ -371,7 +288,6 @@ def __init__(self): model = Model() loader = AutoWeightsLoader(model) - # Weights for the missing layer should not raise loaded = loader.load_weights( [ ("embed.weight", torch.ones(4, 4)), @@ -382,11 +298,6 @@ def __init__(self): self.assertIn("embed.weight", loaded) -# --------------------------------------------------------------------------- -# filter_pp_weights Tests -# --------------------------------------------------------------------------- - - class TestFilterPPWeights(unittest.TestCase): def test_drops_out_of_range_layers(self): weights = [ @@ -412,8 +323,6 @@ def test_passes_non_layer_weights(self): result = list(filter_pp_weights(weights, start_layer=1, end_layer=3)) - # embed_tokens, lm_head, norm pass through (no layer id) - # layer 0 is filtered out names = [name for name, _ in result] self.assertIn("model.embed_tokens.weight", names) self.assertIn("lm_head.weight", names) @@ -433,14 +342,8 @@ def test_empty_range_drops_all_layers(self): self.assertIn("model.embed_tokens.weight", names) -# --------------------------------------------------------------------------- -# WeightRemapRegistry Tests -# --------------------------------------------------------------------------- - - class TestWeightRemapRegistry(unittest.TestCase): def test_registered_model_returns_mapper(self): - # LlamaForCausalLM is registered in auto_loader.py class FakeLlama(nn.Module): pass From 89e4cda686cb286e19e33cc8d22e62b0ddedf185 Mon Sep 17 00:00:00 2001 From: JD-ETH Date: Fri, 10 Jul 2026 19:05:02 +0000 Subject: [PATCH 05/10] Replace unit tests with e2e weight loader verification. 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 --- python/sglang/srt/model_loader/design.md | 2 +- test/manual/test_weight_loader_v2_equiv.py | 143 ++++--- .../test_weight_loader_v2_e2e.py | 86 ++++ .../registered/unit/utils/test_auto_loader.py | 386 ------------------ 4 files changed, 171 insertions(+), 446 deletions(-) create mode 100644 test/registered/model_loading/test_weight_loader_v2_e2e.py delete mode 100644 test/registered/unit/utils/test_auto_loader.py diff --git a/python/sglang/srt/model_loader/design.md b/python/sglang/srt/model_loader/design.md index d5f316081bba..6015d7e95785 100644 --- a/python/sglang/srt/model_loader/design.md +++ b/python/sglang/srt/model_loader/design.md @@ -80,7 +80,7 @@ parent — but are included in the cumulative file count once smoke-tested. | `models/qwen2.py` | v2 gate; `Qwen2Attention` / `Qwen2MLP` submodule loaders | | `models/transformers.py` | already on `AutoWeightsLoader`; document as covered (no v2 gate) | | `environ.py` | `SGLANG_ENABLE_WEIGHT_LOADER_V2 = EnvBool(False)` | -| Tests | `test_auto_loader.py`, `test_auto_loader_moe.py`, `test_weight_loader_v2_equiv.py` | +| Tests | `test_weight_loader_v2_e2e.py` (Qwen2 v1/v2 + transformers backend) | Quant `process_weights_after_loading` remains in `loader.py` (unchanged). diff --git a/test/manual/test_weight_loader_v2_equiv.py b/test/manual/test_weight_loader_v2_equiv.py index cd8aa0856a3a..1cf71c3e9c1a 100644 --- a/test/manual/test_weight_loader_v2_equiv.py +++ b/test/manual/test_weight_loader_v2_equiv.py @@ -11,7 +11,13 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Differential validation: v1 vs v2 weight loaders produce identical state.""" +# 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 @@ -19,66 +25,85 @@ from sglang.srt.environ import envs - -class TestWeightLoaderV2Equivalence(unittest.TestCase): - MODEL_PATH = "Qwen/Qwen2-0.5B" +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_equivalence(self): - from transformers import AutoConfig - - from sglang.srt.model_loader.weight_utils import get_model_weights_iter - from sglang.srt.models.qwen2 import Qwen2ForCausalLM - - config = AutoConfig.from_pretrained(self.MODEL_PATH, trust_remote_code=True) - - with torch.device("cuda"): - model_v1 = Qwen2ForCausalLM(config) - model_v2 = Qwen2ForCausalLM(config) - - torch.manual_seed(0) - for param in (*model_v1.parameters(), *model_v2.parameters()): - param.data.random_() - - weights = list(get_model_weights_iter(self.MODEL_PATH)) - - with envs.SGLANG_ENABLE_WEIGHT_LOADER_V2.override(False): - model_v1.load_weights(iter(weights)) - - with envs.SGLANG_ENABLE_WEIGHT_LOADER_V2.override(True): - loaded_names = model_v2.load_weights(iter(weights)) - - self.assertGreater(len(loaded_names), 0) - - v2_params = dict(model_v2.named_parameters()) - mismatches = [] - for name, p1 in model_v1.named_parameters(): - p2 = v2_params[name] - if not torch.equal(p1.data, p2.data): - max_diff = (p1.data - p2.data).abs().max().item() - mismatches.append(f" {name}: max_diff={max_diff:.6e}") - - if mismatches: - self.fail( - f"Parameter mismatches ({len(mismatches)}):\n" - + "\n".join(mismatches[:20]) - ) - - self.assertTrue(loaded_names.issubset(set(v2_params.keys()))) - - v1_buffers = dict(model_v1.named_buffers()) - v2_buffers = dict(model_v2.named_buffers()) - for name, b1 in v1_buffers.items(): - if name not in v2_buffers: - continue - b2 = v2_buffers[name] - if not torch.equal(b1, b2): - max_diff = (b1 - b2).abs().max().item() - mismatches.append(f" buffer {name}: max_diff={max_diff:.6e}") - - if mismatches: - self.fail( - f"Buffer mismatches ({len(mismatches)}):\n" + "\n".join(mismatches[:20]) + 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, ) 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 000000000000..7ec68c9c3476 --- /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() diff --git a/test/registered/unit/utils/test_auto_loader.py b/test/registered/unit/utils/test_auto_loader.py deleted file mode 100644 index b1fa978b7cca..000000000000 --- a/test/registered/unit/utils/test_auto_loader.py +++ /dev/null @@ -1,386 +0,0 @@ -# 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. -# ============================================================================== -"""Unit tests for the centralized weight loading utilities (auto_loader.py).""" - -from sglang.test.ci.ci_register import register_cpu_ci - -register_cpu_ci(est_time=5, suite="base-a-test-cpu") - -import unittest -from unittest.mock import MagicMock - -import pytest -import torch -from torch import nn -from torch.nn import Parameter - -from sglang.srt.model_loader.auto_loader import ( - STANDARD_GATE_UP_MAPPING, - STANDARD_QKV_MAPPING, - STANDARD_STACKED_MAPPING, - filter_pp_weights, - get_weight_remap, - register_weight_remap, -) -from sglang.srt.models.utils import AutoWeightsLoader, WeightsMapper - - -def _make_param_with_loader(shape): - param = Parameter(torch.zeros(shape)) - param.weight_loader = MagicMock() - return param - - -@pytest.mark.parametrize( - "mapping,source_name,expected_target,expected_shard", - [ - (STANDARD_QKV_MAPPING, "q_proj.weight", "qkv_proj.weight", "q"), - (STANDARD_QKV_MAPPING, "k_proj.weight", "qkv_proj.weight", "k"), - (STANDARD_QKV_MAPPING, "v_proj.weight", "qkv_proj.weight", "v"), - (STANDARD_GATE_UP_MAPPING, "gate_proj.weight", "gate_up_proj.weight", 0), - (STANDARD_GATE_UP_MAPPING, "up_proj.weight", "gate_up_proj.weight", 1), - ], -) -def test_stacked_params_dispatch(mapping, source_name, expected_target, expected_shard): - params = {expected_target: _make_param_with_loader((30, 10))} - tensor = torch.ones(10, 10) - - result = mapping.try_load(source_name, tensor, params) - - assert result == expected_target - params[expected_target].weight_loader.assert_called_once_with( - params[expected_target], tensor, expected_shard - ) - - -def test_stacked_params_dispatch_no_match_returns_none(): - params = {"qkv_proj.weight": _make_param_with_loader((30, 10))} - tensor = torch.ones(10, 10) - - result = STANDARD_QKV_MAPPING.try_load("o_proj.weight", tensor, params) - - assert result is None - params["qkv_proj.weight"].weight_loader.assert_not_called() - - -def test_stacked_params_dispatch_missing_target_param_returns_target_name(): - params = {} - tensor = torch.ones(10) - - result = STANDARD_QKV_MAPPING.try_load("q_proj.bias", tensor, params) - - assert result == "qkv_proj.bias" - - -def test_standard_stacked_mapping_covers_all(): - params = { - "qkv_proj.weight": _make_param_with_loader((30, 10)), - "gate_up_proj.weight": _make_param_with_loader((20, 10)), - } - tensor = torch.ones(10, 10) - - for source, expected_target in [ - ("q_proj.weight", "qkv_proj.weight"), - ("k_proj.weight", "qkv_proj.weight"), - ("v_proj.weight", "qkv_proj.weight"), - ("gate_proj.weight", "gate_up_proj.weight"), - ("up_proj.weight", "gate_up_proj.weight"), - ]: - result = STANDARD_STACKED_MAPPING.try_load(source, tensor, params) - assert result == expected_target, f"Failed for {source}" - - -@pytest.mark.parametrize( - "checkpoint_name,expected_name", - [ - ( - "model.layers.0.self_attn.q_proj.activation_scale", - "model.layers.0.self_attn.q_proj.input_scale", - ), - ( - "model.layers.0.mlp.down_proj.weight_scale_inv", - "model.layers.0.mlp.down_proj.weight_scale", - ), - ( - "model.layers.3.self_attn.k_proj.weight_scale_inv", - "model.layers.3.self_attn.k_proj.weight_scale", - ), - ], -) -def test_llama_remap_registry(checkpoint_name, expected_name): - class FakeLlama(nn.Module): - pass - - FakeLlama.__name__ = "LlamaForCausalLM" - mapper = get_weight_remap(FakeLlama()) - assert mapper is not None - - result = list(mapper.apply([(checkpoint_name, torch.ones(2))])) - assert len(result) == 1 - assert result[0][0] == expected_name - - -def test_weights_mapper_drops_rotary_emb(): - mapper = WeightsMapper(orig_to_new_substr={"rotary_emb": None}) - weights = [ - ("layer.rotary_emb.inv_freq", torch.ones(2)), - ("model.layers.0.self_attn.o_proj.weight", torch.ones(2, 2)), - ] - - result = list(mapper.apply(weights)) - - assert len(result) == 1 - assert result[0][0] == "model.layers.0.self_attn.o_proj.weight" - - -class TestAutoWeightsLoader(unittest.TestCase): - def test_direct_param_load(self): - model = nn.Linear(4, 4, bias=False) - loader = AutoWeightsLoader(model) - tensor = torch.ones(4, 4) - - loaded = loader.load_weights([("weight", tensor)]) - - self.assertIn("weight", loaded) - self.assertTrue(torch.equal(model.weight.data, tensor)) - - def test_nested_module_load(self): - class Model(nn.Module): - def __init__(self): - super().__init__() - self.layer = nn.Linear(4, 4, bias=False) - - model = Model() - loader = AutoWeightsLoader(model) - tensor = torch.ones(4, 4) - - loaded = loader.load_weights([("layer.weight", tensor)]) - - self.assertIn("layer.weight", loaded) - self.assertTrue(torch.equal(model.layer.weight.data, tensor)) - - def test_skip_prefixes(self): - class Model(nn.Module): - def __init__(self): - super().__init__() - self.kept = nn.Linear(4, 4, bias=False) - self.skipped = nn.Linear(4, 4, bias=False) - - model = Model() - loader = AutoWeightsLoader(model, skip_prefixes=["skipped."]) - tensor = torch.ones(4, 4) - - loaded = loader.load_weights( - [ - ("kept.weight", tensor), - ("skipped.weight", tensor * 2), - ] - ) - - self.assertIn("kept.weight", loaded) - self.assertNotIn("skipped.weight", loaded) - self.assertTrue(torch.equal(model.kept.weight.data, tensor)) - - def test_skip_rotary_embeddings(self): - class Model(nn.Module): - def __init__(self): - super().__init__() - self.weight = nn.Parameter(torch.zeros(4)) - - model = Model() - loader = AutoWeightsLoader(model) - - loaded = loader.load_weights( - [ - ("weight", torch.ones(4)), - ("rotary_emb.inv_freq", torch.ones(10)), - ] - ) - - self.assertIn("weight", loaded) - self.assertNotIn("rotary_emb.inv_freq", loaded) - - def test_module_delegation(self): - class ChildModule(nn.Module): - def __init__(self): - super().__init__() - self.param = nn.Parameter(torch.zeros(4)) - - def load_weights(self, weights): - loaded = set() - for name, tensor in weights: - if name == "source_name": - self.param.data.copy_(tensor) - loaded.add("param") - return loaded - - class Model(nn.Module): - def __init__(self): - super().__init__() - self.child = ChildModule() - - model = Model() - loader = AutoWeightsLoader(model) - tensor = torch.ones(4) - - loaded = loader.load_weights([("child.source_name", tensor)]) - - self.assertIn("child.param", loaded) - self.assertTrue(torch.equal(model.child.param.data, tensor)) - - def test_ignore_unexpected_suffix(self): - class Model(nn.Module): - def __init__(self): - super().__init__() - self.layer = nn.Linear(4, 4, bias=False) - - model = Model() - loader = AutoWeightsLoader(model, ignore_unexpected_suffixes=[".bias"]) - - loaded = loader.load_weights( - [ - ("layer.weight", torch.ones(4, 4)), - ("layer.bias", torch.ones(4)), - ] - ) - - self.assertIn("layer.weight", loaded) - - def test_raises_on_truly_unexpected(self): - model = nn.Linear(4, 4, bias=False) - loader = AutoWeightsLoader(model) - - with self.assertRaises(ValueError): - loader.load_weights([("nonexistent_param", torch.ones(4))]) - - def test_mapper_applied(self): - model = nn.Linear(4, 4, bias=False) - mapper = WeightsMapper(orig_to_new_substr={"old_weight": "weight"}) - loader = AutoWeightsLoader(model) - tensor = torch.ones(4, 4) - - loaded = loader.load_weights([("old_weight", tensor)], mapper=mapper) - - self.assertIn("weight", loaded) - self.assertTrue(torch.equal(model.weight.data, tensor)) - - def test_pp_missing_layer_skipped(self): - class PPMissingLayer(nn.Module): - pass - - class Model(nn.Module): - def __init__(self): - super().__init__() - self.layer = PPMissingLayer() - self.embed = nn.Linear(4, 4, bias=False) - - model = Model() - loader = AutoWeightsLoader(model) - - loaded = loader.load_weights( - [ - ("embed.weight", torch.ones(4, 4)), - ("layer.weight", torch.ones(4, 4)), - ] - ) - - self.assertIn("embed.weight", loaded) - - -class TestFilterPPWeights(unittest.TestCase): - def test_drops_out_of_range_layers(self): - weights = [ - ("model.layers.0.self_attn.weight", torch.ones(4)), - ("model.layers.1.self_attn.weight", torch.ones(4)), - ("model.layers.2.self_attn.weight", torch.ones(4)), - ("model.layers.3.self_attn.weight", torch.ones(4)), - ] - - result = list(filter_pp_weights(weights, start_layer=1, end_layer=3)) - - self.assertEqual(len(result), 2) - self.assertEqual(result[0][0], "model.layers.1.self_attn.weight") - self.assertEqual(result[1][0], "model.layers.2.self_attn.weight") - - def test_passes_non_layer_weights(self): - weights = [ - ("model.embed_tokens.weight", torch.ones(4)), - ("lm_head.weight", torch.ones(4)), - ("model.layers.0.self_attn.weight", torch.ones(4)), - ("model.norm.weight", torch.ones(4)), - ] - - result = list(filter_pp_weights(weights, start_layer=1, end_layer=3)) - - names = [name for name, _ in result] - self.assertIn("model.embed_tokens.weight", names) - self.assertIn("lm_head.weight", names) - self.assertIn("model.norm.weight", names) - self.assertNotIn("model.layers.0.self_attn.weight", names) - - def test_empty_range_drops_all_layers(self): - weights = [ - ("model.layers.0.weight", torch.ones(4)), - ("model.embed_tokens.weight", torch.ones(4)), - ] - - result = list(filter_pp_weights(weights, start_layer=5, end_layer=5)) - - names = [name for name, _ in result] - self.assertNotIn("model.layers.0.weight", names) - self.assertIn("model.embed_tokens.weight", names) - - -class TestWeightRemapRegistry(unittest.TestCase): - def test_registered_model_returns_mapper(self): - class FakeLlama(nn.Module): - pass - - FakeLlama.__name__ = "LlamaForCausalLM" - model = FakeLlama() - - mapper = get_weight_remap(model) - - self.assertIsNotNone(mapper) - self.assertIsInstance(mapper, WeightsMapper) - - def test_unknown_model_returns_none(self): - class UnknownModel(nn.Module): - pass - - model = UnknownModel() - - mapper = get_weight_remap(model) - - self.assertIsNone(mapper) - - def test_custom_registration(self): - @register_weight_remap("TestModelXYZ") - def _test_remap(model): - return WeightsMapper(orig_to_new_substr={"foo": "bar"}) - - class TestModelXYZ(nn.Module): - pass - - model = TestModelXYZ() - mapper = get_weight_remap(model) - - self.assertIsNotNone(mapper) - weights = [("prefix.foo.weight", torch.ones(2))] - result = list(mapper.apply(weights)) - self.assertEqual(result[0][0], "prefix.bar.weight") - - -if __name__ == "__main__": - unittest.main() From ab467fac571ef38b733a470cf8633fc36ef81d11 Mon Sep 17 00:00:00 2001 From: JD-ETH Date: Fri, 10 Jul 2026 21:16:07 +0000 Subject: [PATCH 06/10] Document parallel development plan for weight loader PR stack. MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add wave-1 parallelism, sequencing for PR7–PR9, and infra conflict hotspots. Co-authored-by: Cursor --- python/sglang/srt/model_loader/design.md | 70 ++++++++++++++++++++++++ 1 file changed, 70 insertions(+) diff --git a/python/sglang/srt/model_loader/design.md b/python/sglang/srt/model_loader/design.md index 6015d7e95785..67f4a33c2c10 100644 --- a/python/sglang/srt/model_loader/design.md +++ b/python/sglang/srt/model_loader/design.md @@ -305,6 +305,76 @@ welcome; add mapping unit tests and a manual or CI equiv test where a small chec --- +## Parallel development + +### Hard gate + +**PR1 only.** All other PRs need the core stack (`auto_loader.py`, `WeightLoadResult`, dispatch +tables, env gate, `qwen2.py` reference pattern). Downstream branches should stack on PR1 until it +merges. + +### Parallel after PR1 merges + +These tracks can be owned by different people at the same time if they touch **disjoint model +files** and coordinate **shared infra** (`auto_loader.py`, `deepseek_weight_loader.py`): + +| Track | PR | Parallel-safe because | +|-------|-----|----------------------| +| Dense bulk | PR2 (~40 files) | Mechanical `stacked_params_mapping` only; no file overlap with PR3/5/6/7 | +| MoE bulk | PR3 (~22 files) | Standard `ExpertParamsDispatch`; disjoint model list from PR2 | +| MLA / DeepSeek V2 | PR4 (~5 files) | Own mixin + `post_load_weights`; small isolated file set | +| Qwen3.5 fused ckpt | PR5 (`qwen3_5.py`) | Single file; adds `FusedExpertDispatch` in `auto_loader.py` | +| Qwen3-Next hybrid | PR6 (`qwen3_next.py`) | Single file; mamba + MoE submodule loaders | +| DeepSeek V4 | PR7 (`deepseek_v4.py`) | Single file; mostly self-contained bespoke loader | + +**Suggested wave 1 (max parallelism):** PR2 + PR3 + PR4 + PR5 + PR6 in parallel. + +PR7 can start with wave 1 but **reuses PR4 `post_load_weights` protocol** — easiest stacked on +PR4 or merged after PR4. + +### Sequential / later + +| PR | Waits on | Reason | +|----|----------|--------| +| PR8 (MM ~25) | PR2 + PR3 (text trunk); PR5 for VL-MoE fused experts | Trunk uses dense/MoE walker; VL-MoE needs `FusedExpertDispatch` | +| PR9 (wrappers ~35) | Parent model PR per wrapper | No new loader logic; smoke once parent is migrated | + +PR9 is highly parallel **within itself** once parents land (e.g. Qwen wrappers after PR2/5/6, +DeepSeek wrappers after PR4/7). + +### Within-PR parallelism + +| PR | How to split | +|----|--------------| +| PR2 | By family: Llama/Mistral, Gemma, Qwen3, GLM, encoders — ~8–10 independent slices | +| PR3 | One file per model — embarrassingly parallel | +| PR8 | Per tower / per VL model | +| PR9 | Per wrapper file — parallel once parent merges | + +### Infra conflict hotspots + +Serialize or pair-review when multiple tracks touch the same file: + +- `auto_loader.py` — PR5 adds `FusedExpertDispatch`; PR3/PR4 may extend expert/remap behavior +- `deepseek_common/deepseek_weight_loader.py` — PR4 refactors mixin; PR7 may touch shared patterns +- PR2 cherry-picks from brayden branch — one owner for conflict resolution, then parallel model files + +### Suggested merge order + +``` +PR1 (merge) + ├─ PR2 ─┐ + ├─ PR3 ─┼─ wave 1 (parallel) + ├─ PR4 ─┤ + ├─ PR5 ─┤ + └─ PR6 ─┘ + ├─ PR7 (stack on PR4) + ├─ PR8 (after PR2 + PR3; VL-MoE after PR5) + └─ PR9 (sharded by parent; ongoing as parents land) +``` + +--- + ## Checklist - [ ] PR1: infra + post_load protocol + `qwen2.py` + `transformers.py` (walker) From 6d1709dda6502d053c6d21ed18538455ad536bef Mon Sep 17 00:00:00 2001 From: JD-ETH Date: Mon, 13 Jul 2026 17:26:02 +0000 Subject: [PATCH 07/10] Move weight loader v2 design doc to issue #31051. Remove design.md from the PR; point auto_loader docstring at the migration tracker issue linked from RFC #24703. Co-authored-by: Cursor --- python/sglang/srt/model_loader/auto_loader.py | 4 +- python/sglang/srt/model_loader/design.md | 395 ------------------ 2 files changed, 2 insertions(+), 397 deletions(-) delete mode 100644 python/sglang/srt/model_loader/design.md diff --git a/python/sglang/srt/model_loader/auto_loader.py b/python/sglang/srt/model_loader/auto_loader.py index 7bcb8e2eb2a6..259c07dd4961 100644 --- a/python/sglang/srt/model_loader/auto_loader.py +++ b/python/sglang/srt/model_loader/auto_loader.py @@ -20,11 +20,11 @@ - RemapRegistry: architecture-specific name remap registration. - Re-exports of AutoWeightsLoader and WeightsMapper from models/utils.py. -Load / post-load split (PR1 protocol, see design.md): +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: design.md in this directory. +Migration plan: https://github.com/sgl-project/sglang/issues/31051 (RFC #24703). """ from __future__ import annotations diff --git a/python/sglang/srt/model_loader/design.md b/python/sglang/srt/model_loader/design.md deleted file mode 100644 index 67f4a33c2c10..000000000000 --- a/python/sglang/srt/model_loader/design.md +++ /dev/null @@ -1,395 +0,0 @@ -# Weight Loader v2 — design and migration plan - -Tracking issue: [RFC #24703](https://github.com/sgl-project/sglang/issues/24703) -Demo PR: [#28671](https://github.com/sgl-project/sglang/pull/28671) - ---- - -## Why this design - -**Complete loading, no stragglers.** Every checkpoint tensor that enters the loader and is not -explicitly skipped must produce a `WeightLoadRecord`, or `verify_complete` raises -`IncompleteWeightLoadError`. That replaces silent `logger.warning` drops and ad-hoc expert-shard -checks scattered across individual model files. - -**Clean, visible weight mapping.** Each load is recorded as -`(checkpoint_name → runtime_name, shard_id, expert_id)`. Stacked qkv/gate_up and MoE expert -shards are first-class — not buried inside 70-line per-model loops. Mapping tables live in -`StackedParamsDispatch`, `ExpertParamsDispatch`, and `RemapRegistry` instead of being -copy-pasted 116 times across `srt/models/`. - -**Native weight-loading completeness check.** v2 loads call `verify_complete` to assert every -expected runtime tensor was updated — no silent skips, no stragglers. A mergeable -`WeightLoadResult` with per-shard records supports this check across single or batched -`load_weights` calls. The load / `post_load_weights` split applies completeness verification -after placement and before optional GPU-specific derivation (MLA `w_kc`/`w_vc`, etc.). - -All migrations stay behind `SGLANG_ENABLE_WEIGHT_LOADER_V2` (opt-in, default off). Legacy -`load_weights` paths remain until explicitly removed per model. - ---- - -## Migration overview (~180 `load_weights` files) - -Cumulative auto-loader coverage by PR (unique model implementation files, ~180 total): - -| PR | Adds (files) | Cumulative | Coverage | -|----|-------------:|-----------:|---------:| -| **PR1** | 2 (`qwen2.py`, `transformers.py`) | 2 | 1% | -| **PR2** | ~40 | ~42 | 23% | -| **PR3** | ~22 | ~64 | 36% | -| **PR4** | ~5 | ~69 | 38% | -| **PR5** | 1 | ~70 | 39% | -| **PR6** | 1 | ~71 | 39% | -| **PR7** | 1 | ~72 | 40% | -| **PR8** | ~25 | ~97 | 54% | -| **PR9** | ~35 wrappers | ~132 | 73% | - -`transformers.py` uses `AutoWeightsLoader` today (no v2 gate); counted in PR1 as already on the -walker stack. PR9 wrapper files add no new loader logic — they inherit v2 from a migrated -parent — but are included in the cumulative file count once smoke-tested. - -### Protocol extensions (introduced across the PR stack) - -| Extension | PR | Purpose | -|-----------|-----|---------| -| `post_load_weights` split | PR1 / PR4 | Separate tensor load from GPU derivations | -| Shared-expert preload hook | PR4 | Rename `shared_experts` → routed expert slot before dispatch | -| MLA `RemapRegistry` entries | PR4 | `fused_qkv_a_proj_with_mqa`, scale remaps | -| `FusedExpertDispatch` | PR5 | One fused ckpt key → multiple expert shard records | -| Packed-shard fan-out | PR5 | One ckpt tensor → multiple column shard records | -| Hybrid mixer submodule loaders | PR6 | Mamba + MoE in one decoder block | -| Custom loader decomposition | PR7 | Large bespoke loops → submodule overrides | -| Vision / audio tower `load_weights` | PR8 | Walker delegates to per-tower loaders | - ---- - -## PR1 — Infrastructure + first model (this PR) - -**Addresses:** Core abstractions and the reference dense migration pattern. - -**Coverage:** +2 files (`qwen2.py`, `transformers.py` on walker stack) (~1%). - -**Work:** - -| Area | Deliverable | -|------|-------------| -| `auto_loader.py` | `WeightLoadRecord`, `WeightLoadResult`, `StackedParamsDispatch`, `ExpertParamsDispatch`, `RemapRegistry`, `verify_complete` | -| `models/utils.py` | `AutoWeightsLoader` returns `WeightLoadResult` | -| Load protocol | `load_weights(..., run_post_load=True)` + default no-op `post_load_weights()` | -| `models/qwen2.py` | v2 gate; `Qwen2Attention` / `Qwen2MLP` submodule loaders | -| `models/transformers.py` | already on `AutoWeightsLoader`; document as covered (no v2 gate) | -| `environ.py` | `SGLANG_ENABLE_WEIGHT_LOADER_V2 = EnvBool(False)` | -| Tests | `test_weight_loader_v2_e2e.py` (Qwen2 v1/v2 + transformers backend) | - -Quant `process_weights_after_loading` remains in `loader.py` (unchanged). - ---- - -## PR2 — Standard dense models (~41 files) - -**Addresses:** Models with `stacked_params_mapping` only (qkv + gate_up), no expert mapping, -no MLA mixin. Mechanical migration: env gate, submodule loaders, optional `RemapRegistry`. - -**Coverage:** +~40 files (~23% cumulative). - -**Work per file:** Rename loop → `_legacy_load_weights`; add `_load_weights_v2` using -`AutoWeightsLoader` + `StackedParamsDispatch` on attention/MLP submodules; register FP8/kv -remaps where needed. - -**Models:** - -| File | Notes | -|------|-------| -| `qwen2.py` | done in PR1 | -| `qwen3.py` | + `maybe_remap_kv_scale` | -| `qwen3_classification.py` | inherits Qwen3 | -| `llama.py` | FP8 suffix + kv scale via `RemapRegistry` | -| `mistral.py` | delegates → Llama | -| `glm4.py` | dot-prefix stacked | -| `gemma.py`, `gemma2.py`, `gemma3_causal.py`, `gemma3n_causal.py` | gemma2/3 + kv remap | -| `olmo.py`, `olmo2.py`, `internlm2.py`, `granite.py` | standard | -| `stablelm.py`, `starcoder2.py`, `opt.py`, `orion.py` | standard | -| `baichuan.py`, `xverse.py`, `exaone.py`, `exaone4.py` | standard / remap | -| `commandr.py`, `apertus.py`, `arcee.py`, `nemotron_nas.py` | + remap | -| `falcon_h1.py`, `lfm2.py`, `iquest_loopcoder.py`, `jet_nemotron.py` | standard | -| `mimo.py`, `gpt_j.py`, `sdar.py`, `dflash.py`, `afmoe.py` | standard / remap | -| `bert.py`, `roberta.py`, `clip.py`, `whisper.py`, `siglip2.py` | encoders | -| `llama_embedding.py`, `torch_native_llama.py` | variants | - -Cherry-pick walker diffs from -[brayden/weight-loader-refactor-part-1](https://github.com/sgl-project/sglang/compare/main...brayden/weight-loader-refactor-part-1) -where possible. - -**Tests:** parametrized CPU mapping tests; spot-check GPU equiv (`Llama-3.2-1B`, `Qwen3-1.7B-FP8`). - ---- - -## PR3 — Standard MoE models (~22 files) - -**Addresses:** Models using `FusedMoE.make_expert_params_mapping` with dense stacked layers. - -**Coverage:** +~22 files (~36% cumulative). - -**Work per file:** `ExpertParamsDispatch` on `FusedMoE` submodule; dense submodules as PR2; -`verify_expert_shards_complete` where all routed expert shards are mandatory (e.g. `laguna.py`). - -**Models:** - -| File | Notes | -|------|-------| -| `qwen2_moe.py`, `qwen3_moe.py` | standard MoE | -| `mixtral.py`, `glm4_moe.py`, `olmoe.py`, `phimoe.py` | standard | -| `cohere2_moe.py`, `exaone_moe.py`, `lfm2_moe.py` | standard | -| `hunyuan.py`, `hunyuan_v3.py`, `bailing_moe.py` | standard | -| `laguna.py` | + expert-shard completeness assert | -| `ernie4.py`, `grok.py`, `kimi_linear.py`, `llada2.py` | standard | -| `minimax_m2.py`, `sdar_moe.py`, `sarvam_moe.py`, `dbrx.py` | standard | -| `xverse_moe.py`, `granitemoe.py` | standard / delegates Mixtral | - -**Tests:** `Qwen/Qwen1.5-MoE-A2.7B` equiv; extend `test_auto_loader_moe.py`. - ---- - -## PR4 — MLA mixin models (~5 files) - -**Addresses:** `DeepseekV2WeightLoaderMixin` / `do_load_weights` — shared expert preload, MLA -name remaps, and `post_load_weights` for `w_kc`/`w_vc` derivation. - -**Coverage:** +~5 files (~38% cumulative). - -**Work:** Replace mixin loop with walker + submodule loaders; wire shared-expert preload hook; -split `post_load_weights` from tensor load (PR1 protocol); shrink `deepseek_weight_loader.py`. - -**Models:** - -| File | Notes | -|------|-------| -| `deepseek_v2.py` | V2, V3, V32 | -| `deepseek_common/deepseek_weight_loader.py` | mixin refactor | -| `glm4_moe_lite.py` | GLM-4.7 | -| `bailing_moe_linear.py` | linear MoE variant | -| `kimi_k25_eagle3.py` | eagle MLA draft | - -**Tests:** `deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct`; GLM-4.7 mapping units. - ---- - -## PR5 — `qwen3_5.py` - -**Addresses:** Fused MoE checkpoint format (`gate_up_proj` → w1 + w3 shards) and packed linear -shard fan-out. - -**Coverage:** +1 file (~39% cumulative). - -**Work:** Add `FusedExpertDispatch`; fan-out multi-shard records from single ckpt tensors; -migrate `Qwen3_5MoeForConditionalGeneration`, `Qwen3_5ForConditionalGeneration`. - -**Tests:** `Qwen/Qwen3.5-4B`; fused-expert unit tests. - ---- - -## PR6 — `qwen3_next.py` - -**Addresses:** Hybrid mamba + MoE decoder — mixer and expert paths need separate submodule -`load_weights`. - -**Coverage:** +1 file (~39% cumulative). - -**Work:** Submodule overrides on mamba mixer and MoE block; standard dispatch elsewhere. - -**Tests:** mapping unit tests (full-model GPU equiv impractical at default sizes). - ---- - -## PR7 — `deepseek_v4.py` - -**Addresses:** Large bespoke loader (~335 lines) with custom expert routing and `post_load_weights`. - -**Coverage:** +1 file (~40% cumulative). - -**Work:** Decompose into submodule loaders; reuse PR4 post_load protocol. - -**Tests:** unit mapping tests; targeted manual validation. - ---- - -## PR8 — Multimodal models (~25 files) - -**Addresses:** Vision, audio, and OCR tower weights routed alongside text trunk. - -**Coverage:** +~25 files (~54% cumulative). - -**Work:** Text trunk uses PR2/PR3 walker; add `VisionTower.load_weights` / -`AudioTower.load_weights` submodule delegation; `skip_substrs` for optional towers. Reuse -`FusedExpertDispatch` from PR5 where VL-MoE uses fused expert ckpts. - -**Models (representative):** `qwen2_vl.py`, `qwen2_5_vl.py`, `qwen3_vl.py`, `qwen3_vl_moe.py`, -`qwen3_omni_moe.py`, `qwen2_audio.py`, `qwen3_asr.py`, `deepseek_vl2.py`, `deepseek_ocr.py`, -`deepseek_janus_pro.py`, `glm4v.py`, `glm4v_moe.py`, `glm_ocr.py`, `glmasr.py`, `minicpmv.py`, -`internvl.py`, `kimi_vl.py`, `moss_vl.py`, `gemma3_mm.py`, `gemma4_mm.py`, `pixtral.py`, … - -**Tests:** `Qwen/Qwen2.5-VL-3B-Instruct`, `Qwen/Qwen3-VL-2B-Thinking`; per-tower CPU units. - ---- - -## PR9 — Wrapper models (~35 files) - -**Addresses:** Thin classes that only call `super().load_weights` (eagle, MTP, NextN, -classification, reward, delegates). - -**Coverage:** +~35 wrapper files (~73% cumulative). No new loader implementation — inherits -v2 from parent. Smoke-test that `super().load_weights` forwards `WeightLoadResult` and -`run_post_load` correctly; fix NextN `__init__` bypass mapping gaps. - -**Work:** Ensure mappings are available when `__init__` is bypassed (NextN pattern); forward -`run_post_load` to parent; smoke tests only. - -**Models (representative):** `qwen2_eagle.py`, `qwen3_moe_mtp.py`, `qwen3_5_mtp.py`, -`qwen3_next_mtp.py`, `deepseek_nextn.py`, `deepseek_v4_nextn.py`, `llava.py`, `granitemoe.py`, -`mistral.py`, `qwen2_rm.py`, … - ---- - -## Not supported today (open to community contribution) - -The PR1–PR9 plan does **not** include these models. They stay on legacy `load_weights` until -someone contributes a dedicated migration. v2 remains opt-in via `SGLANG_ENABLE_WEIGHT_LOADER_V2`. - -### Legacy architectures - -| File | Reason | -|------|--------| -| `qwen.py` | Superseded by Qwen2+; hand-rolled loader | -| `deepseek.py` | Superseded by DeepSeek V2+; hand-rolled loader | - -### Custom loaders (no `stacked_params_mapping` pattern) - -| File | Reason | -|------|--------| -| `phi.py` | `packed_modules_mapping` driven | -| `phi3_small.py` | Custom loader loop | -| `solar.py` | `packed_modules_mapping` driven | -| `gpt2.py` | Direct param iteration | -| `gpt_bigcode.py` | Direct param iteration | -| `chatglm.py` | Custom GLM QKV routing | -| `minicpm.py` | Manual w1/w2/w3 expert table | -| `minicpm3.py` | Manual expert mapping | -| `parakeet.py` | Custom loader | -| `persimmon.py` | Custom loader | - -### Bespoke loaders (need dedicated design per model) - -| File | Reason | -|------|--------| -| `step3p5.py` | Custom expert name matching (~146L loader) | -| `step3_vl.py` | Custom expert matching + vision | -| `step3p5_mtp.py` | MTP variant of step3 loader | -| `step3_vl_10b.py` | VL delegate + step3 patterns | -| `step3p7.py` | Step family bespoke routing | -| `nemotron_h.py` | Non-standard expert `weight_loader` signature | -| `nemotron_h_mtp.py` | MTP wrapper on nemotron_h | -| `zaya.py` | MOD routing / custom MoE | -| `longcat_flash.py` | Hybrid loader (~180L) | -| `longcat_flash_nextn.py` | NextN + longcat patterns | -| `mimo_v2.py` | Custom qkv routing (~198L loader) | -| `mimo_v2_nextn.py` | MTP on mimo_v2 | -| `mimo_v2_asr.py` | ASR delegate on mimo_v2 | -| `gpt_oss.py` | Unique quant routing | -| `gemma4_causal.py` | Fused expert ckpt (needs `FusedExpertDispatch` beyond PR5 scope) | -| `granitemoehybrid.py` | Mamba + MoE hybrid (~150L loader) | -| `mindspore.py` | MindSpore cell loader path | - -Pull requests that migrate any of the above to the v2 stack (walker + dispatch + records) are -welcome; add mapping unit tests and a manual or CI equiv test where a small checkpoint exists. - ---- - -## Parallel development - -### Hard gate - -**PR1 only.** All other PRs need the core stack (`auto_loader.py`, `WeightLoadResult`, dispatch -tables, env gate, `qwen2.py` reference pattern). Downstream branches should stack on PR1 until it -merges. - -### Parallel after PR1 merges - -These tracks can be owned by different people at the same time if they touch **disjoint model -files** and coordinate **shared infra** (`auto_loader.py`, `deepseek_weight_loader.py`): - -| Track | PR | Parallel-safe because | -|-------|-----|----------------------| -| Dense bulk | PR2 (~40 files) | Mechanical `stacked_params_mapping` only; no file overlap with PR3/5/6/7 | -| MoE bulk | PR3 (~22 files) | Standard `ExpertParamsDispatch`; disjoint model list from PR2 | -| MLA / DeepSeek V2 | PR4 (~5 files) | Own mixin + `post_load_weights`; small isolated file set | -| Qwen3.5 fused ckpt | PR5 (`qwen3_5.py`) | Single file; adds `FusedExpertDispatch` in `auto_loader.py` | -| Qwen3-Next hybrid | PR6 (`qwen3_next.py`) | Single file; mamba + MoE submodule loaders | -| DeepSeek V4 | PR7 (`deepseek_v4.py`) | Single file; mostly self-contained bespoke loader | - -**Suggested wave 1 (max parallelism):** PR2 + PR3 + PR4 + PR5 + PR6 in parallel. - -PR7 can start with wave 1 but **reuses PR4 `post_load_weights` protocol** — easiest stacked on -PR4 or merged after PR4. - -### Sequential / later - -| PR | Waits on | Reason | -|----|----------|--------| -| PR8 (MM ~25) | PR2 + PR3 (text trunk); PR5 for VL-MoE fused experts | Trunk uses dense/MoE walker; VL-MoE needs `FusedExpertDispatch` | -| PR9 (wrappers ~35) | Parent model PR per wrapper | No new loader logic; smoke once parent is migrated | - -PR9 is highly parallel **within itself** once parents land (e.g. Qwen wrappers after PR2/5/6, -DeepSeek wrappers after PR4/7). - -### Within-PR parallelism - -| PR | How to split | -|----|--------------| -| PR2 | By family: Llama/Mistral, Gemma, Qwen3, GLM, encoders — ~8–10 independent slices | -| PR3 | One file per model — embarrassingly parallel | -| PR8 | Per tower / per VL model | -| PR9 | Per wrapper file — parallel once parent merges | - -### Infra conflict hotspots - -Serialize or pair-review when multiple tracks touch the same file: - -- `auto_loader.py` — PR5 adds `FusedExpertDispatch`; PR3/PR4 may extend expert/remap behavior -- `deepseek_common/deepseek_weight_loader.py` — PR4 refactors mixin; PR7 may touch shared patterns -- PR2 cherry-picks from brayden branch — one owner for conflict resolution, then parallel model files - -### Suggested merge order - -``` -PR1 (merge) - ├─ PR2 ─┐ - ├─ PR3 ─┼─ wave 1 (parallel) - ├─ PR4 ─┤ - ├─ PR5 ─┤ - └─ PR6 ─┘ - ├─ PR7 (stack on PR4) - ├─ PR8 (after PR2 + PR3; VL-MoE after PR5) - └─ PR9 (sharded by parent; ongoing as parents land) -``` - ---- - -## Checklist - -- [ ] PR1: infra + post_load protocol + `qwen2.py` + `transformers.py` (walker) -- [ ] PR2: standard dense (~41 files) -- [ ] PR3: standard MoE (~22 files) -- [ ] PR4: MLA mixin (~5 files) -- [ ] PR5: `qwen3_5.py` -- [ ] PR6: `qwen3_next.py` -- [ ] PR7: `deepseek_v4.py` -- [ ] PR8: multimodal (~25 files) -- [ ] PR9: wrappers (~35 files) - ---- - -## References - -- [brayden/weight-loader-refactor-part-1](https://github.com/sgl-project/sglang/compare/main...brayden/weight-loader-refactor-part-1) — cherry-pick dense walker diffs into PR2 -- [PR #24256](https://github.com/sgl-project/sglang/pull/24256) — checksum for production equiv validation From f64a164a7a4234fd343db7324d7d902c6df561a6 Mon Sep 17 00:00:00 2001 From: JD-ETH Date: Wed, 15 Jul 2026 02:37:42 +0000 Subject: [PATCH 08/10] feat(model_loader): v2 dense loaders for PR2 (#31051) Add load_with_stacked_dispatch helper and migrate P0 dense models to env-gated AutoWeightsLoader v2 paths with legacy preservation and tests. Co-authored-by: Cursor --- python/sglang/srt/model_loader/auto_loader.py | 26 ++++ python/sglang/srt/models/gemma2.py | 44 ++++++ python/sglang/srt/models/glm4.py | 58 ++++++++ python/sglang/srt/models/granite.py | 50 +++++++ python/sglang/srt/models/internlm2.py | 61 +++++++++ python/sglang/srt/models/olmo.py | 49 +++++++ python/sglang/srt/models/olmo2.py | 49 +++++++ python/sglang/srt/models/qwen3.py | 70 ++++++++++ .../sglang/srt/models/qwen3_classification.py | 23 +++- test/manual/test_weight_loader_v2_equiv.py | 29 +++- .../test_weight_loader_v2_e2e.py | 23 ++++ .../test_stacked_params_dispatch.py | 126 ++++++++++++++++++ 12 files changed, 606 insertions(+), 2 deletions(-) create mode 100644 test/registered/unit/model_loader/test_stacked_params_dispatch.py diff --git a/python/sglang/srt/model_loader/auto_loader.py b/python/sglang/srt/model_loader/auto_loader.py index 259c07dd4961..81e7b9f0da38 100644 --- a/python/sglang/srt/model_loader/auto_loader.py +++ b/python/sglang/srt/model_loader/auto_loader.py @@ -38,6 +38,7 @@ from torch.nn import Parameter from sglang.srt.layers.utils.common import get_layer_id +from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.models.utils import AutoWeightsLoader, WeightsMapper __all__ = [ @@ -48,6 +49,7 @@ "STANDARD_GATE_UP_MAPPING", "STANDARD_STACKED_MAPPING", "LLAMA_STACKED_MAPPING", + "load_with_stacked_dispatch", "filter_pp_weights", "register_weight_remap", "get_weight_remap", @@ -148,6 +150,30 @@ def try_load( ) +def load_with_stacked_dispatch( + module: nn.Module, + weights: Iterable[tuple[str, torch.Tensor]], + mapping: StackedParamsDispatch, + *, + ignore_unexpected_suffixes: tuple[str, ...] = (".bias", ".kv_scale"), +) -> set[str]: + """Load submodule weights via stacked dispatch, then direct param loaders.""" + loaded: set[str] = set() + params_dict = dict(module.named_parameters()) + for name, tensor in weights: + target = 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) + elif not any(name.endswith(suffix) for suffix in ignore_unexpected_suffixes): + pass + return loaded + + # --------------------------------------------------------------------------- # Pipeline Parallel Weight Filter # --------------------------------------------------------------------------- diff --git a/python/sglang/srt/models/gemma2.py b/python/sglang/srt/models/gemma2.py index 4c0942839726..ec13f6f90aae 100644 --- a/python/sglang/srt/models/gemma2.py +++ b/python/sglang/srt/models/gemma2.py @@ -92,6 +92,14 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: 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, + load_with_stacked_dispatch, + ) + + return load_with_stacked_dispatch(self, weights, STANDARD_GATE_UP_MAPPING) + class Gemma2Attention(nn.Module): def __init__( @@ -200,6 +208,14 @@ def forward( 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, + load_with_stacked_dispatch, + ) + + return load_with_stacked_dispatch(self, weights, STANDARD_QKV_MAPPING) + class Gemma2DecoderLayer(nn.Module): def __init__( @@ -458,6 +474,13 @@ def get_attention_sliding_window_size(self): return get_attention_sliding_window_size(self.config) 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"), @@ -498,5 +521,26 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): weight_loader(param, loaded_weight) loaded_params.add(name) + def _load_weights_v2(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> set[str]: + from sglang.srt.model_loader.auto_loader import AutoWeightsLoader + + params_dict = dict(self.named_parameters()) + + def _prepare( + src: Iterable[Tuple[str, torch.Tensor]], + ) -> Iterable[Tuple[str, torch.Tensor]]: + for name, loaded_weight in src: + remapped = maybe_remap_kv_scale_name(name, params_dict) + if remapped is None: + continue + yield remapped, loaded_weight + + loader = AutoWeightsLoader( + self, + skip_prefixes=["lm_head."], + ignore_unexpected_suffixes=[".bias", ".kv_scale"], + ) + return loader.load_weights(_prepare(weights)) + EntryClass = Gemma2ForCausalLM diff --git a/python/sglang/srt/models/glm4.py b/python/sglang/srt/models/glm4.py index 016941b4b6c6..7e777ffe9d7f 100644 --- a/python/sglang/srt/models/glm4.py +++ b/python/sglang/srt/models/glm4.py @@ -105,6 +105,14 @@ def forward( ) 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, + load_with_stacked_dispatch, + ) + + return load_with_stacked_dispatch(self, weights, STANDARD_GATE_UP_MAPPING) + class Glm4Attention(nn.Module): def __init__( @@ -200,6 +208,14 @@ def forward( 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, + load_with_stacked_dispatch, + ) + + return load_with_stacked_dispatch(self, weights, STANDARD_QKV_MAPPING) + class Glm4DecoderLayer(nn.Module): """A single transformer layer. @@ -557,6 +573,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"), @@ -619,6 +642,41 @@ 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]: + from sglang.srt.model_loader.auto_loader import ( + AutoWeightsLoader, + filter_pp_weights, + ) + + 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"], + ignore_unexpected_suffixes=[".bias", ".kv_scale"], + ) + loaded = loader.load_weights(weights) + + 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/python/sglang/srt/models/granite.py b/python/sglang/srt/models/granite.py index eb9d2aa7f2b4..4a66a2849b0c 100644 --- a/python/sglang/srt/models/granite.py +++ b/python/sglang/srt/models/granite.py @@ -88,6 +88,14 @@ def forward(self, x): 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, + load_with_stacked_dispatch, + ) + + return load_with_stacked_dispatch(self, weights, STANDARD_GATE_UP_MAPPING) + class GraniteAttention(nn.Module): def __init__( @@ -178,6 +186,14 @@ def forward( 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, + load_with_stacked_dispatch, + ) + + return load_with_stacked_dispatch(self, weights, STANDARD_QKV_MAPPING) + class GraniteDecoderLayer(nn.Module): def __init__( @@ -380,6 +396,13 @@ def get_num_params(self): 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"), @@ -431,6 +454,33 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight) + def _load_weights_v2(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> set[str]: + from sglang.srt.model_loader.auto_loader import AutoWeightsLoader + + 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"], + ) + loaded = loader.load_weights(weights) + + 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/internlm2.py b/python/sglang/srt/models/internlm2.py index 6a276ef0a7f8..0cf0ef24b380 100644 --- a/python/sglang/srt/models/internlm2.py +++ b/python/sglang/srt/models/internlm2.py @@ -81,10 +81,25 @@ def forward(self, x): x, _ = self.w2(x) return x + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> set[str]: + from sglang.srt.model_loader.auto_loader import ( + StackedParamsDispatch, + load_with_stacked_dispatch, + ) + + mapping = StackedParamsDispatch( + mappings=( + ("gate_up_proj", "w1", 0), + ("gate_up_proj", "w3", 1), + ) + ) + return load_with_stacked_dispatch(self, weights, mapping) + class InternLM2Attention(nn.Module): def __init__( self, + config: PretrainedConfig, hidden_size: int, num_heads: int, num_kv_heads: int, @@ -96,6 +111,7 @@ def __init__( prefix: str = "", ) -> None: super().__init__() + self.config = config self.hidden_size = hidden_size tp_size = get_tensor_model_parallel_world_size() self.total_num_heads = num_heads @@ -165,6 +181,34 @@ def forward( output, _ = self.wo(attn_output) return output + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> set[str]: + loaded: set[str] = set() + params_dict = dict(self.named_parameters()) + config = self.config + for name, tensor in weights: + if "wqkv" in name and name in params_dict: + param = params_dict[name] + kv_groups = config.num_attention_heads // config.num_key_value_heads + head_dim = config.hidden_size // config.num_attention_heads + loaded_weight = tensor.view( + -1, 2 + kv_groups, head_dim, tensor.shape[-1] + ) + wq, wk, wv = torch.split(loaded_weight, [kv_groups, 1, 1], dim=1) + wq = wq.reshape(-1, wq.shape[-1]) + wk = wk.reshape(-1, wk.shape[-1]) + wv = wv.reshape(-1, wv.shape[-1]) + weight_loader = param.weight_loader + weight_loader(param, wq, "q") + weight_loader(param, wk, "k") + weight_loader(param, wv, "v") + loaded.add(name) + 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 InternLMDecoderLayer(nn.Module): def __init__( @@ -180,6 +224,7 @@ def __init__( rope_scaling = getattr(config, "rope_scaling", None) max_position_embeddings = getattr(config, "max_position_embeddings", 8192) self.attention = InternLM2Attention( + config=config, hidden_size=self.hidden_size, num_heads=config.num_attention_heads, num_kv_heads=config.num_key_value_heads, @@ -310,6 +355,13 @@ def forward( ) 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) ("gate_up_proj", "w1", 0), @@ -356,5 +408,14 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): ) weight_loader(param, loaded_weight) + def _load_weights_v2(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> set[str]: + from sglang.srt.model_loader.auto_loader import AutoWeightsLoader + + loader = AutoWeightsLoader( + self, + ignore_unexpected_suffixes=[".bias", ".kv_scale"], + ) + return loader.load_weights(weights) + EntryClass = InternLM2ForCausalLM diff --git a/python/sglang/srt/models/olmo.py b/python/sglang/srt/models/olmo.py index 5f51e229af39..88d6f29649eb 100644 --- a/python/sglang/srt/models/olmo.py +++ b/python/sglang/srt/models/olmo.py @@ -124,6 +124,14 @@ def forward( 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, + load_with_stacked_dispatch, + ) + + return load_with_stacked_dispatch(self, weights, STANDARD_QKV_MAPPING) + class OlmoMLP(nn.Module): """ @@ -173,6 +181,14 @@ def forward( 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, + load_with_stacked_dispatch, + ) + + return load_with_stacked_dispatch(self, weights, STANDARD_GATE_UP_MAPPING) + class OlmoDecoderLayer(nn.Module): """ @@ -341,6 +357,13 @@ def forward( ) 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"), @@ -378,5 +401,31 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight) + def _load_weights_v2(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> set[str]: + from sglang.srt.model_loader.auto_loader import AutoWeightsLoader + + skip_prefixes = [] + if self.config.tie_word_embeddings: + skip_prefixes.append("lm_head.") + + loader = AutoWeightsLoader( + self, + skip_prefixes=skip_prefixes, + ignore_unexpected_suffixes=[".bias", ".kv_scale"], + ) + loaded = loader.load_weights(weights) + + if self.config.tie_word_embeddings: + params_dict = dict(self.named_parameters(remove_duplicate=False)) + 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 + EntryClass = OlmoForCausalLM diff --git a/python/sglang/srt/models/olmo2.py b/python/sglang/srt/models/olmo2.py index 91dfd12cd2ed..1ee4b9a316ca 100644 --- a/python/sglang/srt/models/olmo2.py +++ b/python/sglang/srt/models/olmo2.py @@ -212,6 +212,14 @@ def forward( 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, + load_with_stacked_dispatch, + ) + + return load_with_stacked_dispatch(self, weights, STANDARD_QKV_MAPPING) + class Olmo2MLP(nn.Module): """ @@ -261,6 +269,14 @@ def forward( 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, + load_with_stacked_dispatch, + ) + + return load_with_stacked_dispatch(self, weights, STANDARD_GATE_UP_MAPPING) + class Olmo2DecoderLayer(nn.Module): """ @@ -442,6 +458,13 @@ def forward( ) 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"), @@ -482,5 +505,31 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight) + def _load_weights_v2(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> set[str]: + from sglang.srt.model_loader.auto_loader import AutoWeightsLoader + + skip_prefixes = [] + if self.config.tie_word_embeddings: + skip_prefixes.append("lm_head.") + + loader = AutoWeightsLoader( + self, + skip_prefixes=skip_prefixes, + ignore_unexpected_suffixes=[".bias", ".kv_scale"], + ) + loaded = loader.load_weights(weights) + + if self.config.tie_word_embeddings: + params_dict = dict(self.named_parameters(remove_duplicate=False)) + 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 + EntryClass = Olmo2ForCausalLM diff --git a/python/sglang/srt/models/qwen3.py b/python/sglang/srt/models/qwen3.py index 14fb21dfc7eb..55daec523d5e 100644 --- a/python/sglang/srt/models/qwen3.py +++ b/python/sglang/srt/models/qwen3.py @@ -312,6 +312,32 @@ def forward( 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, + load_with_stacked_dispatch, + ) + + return load_with_stacked_dispatch(self, weights, STANDARD_QKV_MAPPING) + + +def _prepare_qwen3_checkpoint_weights( + weights: Iterable[Tuple[str, torch.Tensor]], + params_dict: dict[str, torch.nn.Parameter], +) -> Iterable[Tuple[str, torch.Tensor]]: + for name, loaded_weight in weights: + if not name.startswith("model.") and ( + name.startswith("layers.") + or name.startswith("embed_tokens.") + or name.startswith("norm.") + ): + name = add_prefix(name, "model") + if "scale" in name: + name = maybe_remap_kv_scale_name(name, params_dict) + if name is None: + continue + yield name, loaded_weight + class Qwen3DecoderLayer(nn.Module): def __init__( @@ -600,6 +626,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"), @@ -675,6 +708,43 @@ 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]: + from sglang.srt.model_loader.auto_loader import ( + AutoWeightsLoader, + filter_pp_weights, + ) + + params_dict = dict(self.named_parameters()) + weights = _prepare_qwen3_checkpoint_weights(weights, params_dict) + + 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"], + ) + loaded = loader.load_weights(weights) + + if self.config.tie_word_embeddings and self.pp_group.is_last_rank: + 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/python/sglang/srt/models/qwen3_classification.py b/python/sglang/srt/models/qwen3_classification.py index 352f7c7d8434..9adfff0d6d75 100644 --- a/python/sglang/srt/models/qwen3_classification.py +++ b/python/sglang/srt/models/qwen3_classification.py @@ -28,7 +28,7 @@ from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_loader.weight_utils import default_weight_loader -from sglang.srt.models.qwen3 import Qwen3Model +from sglang.srt.models.qwen3 import Qwen3Model, _prepare_qwen3_checkpoint_weights from sglang.srt.utils import add_prefix logger = logging.getLogger(__name__) @@ -75,6 +75,13 @@ def forward( ) 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"), @@ -124,6 +131,20 @@ 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]: + from sglang.srt.model_loader.auto_loader import AutoWeightsLoader + + params_dict = dict(self.named_parameters()) + weights = _prepare_qwen3_checkpoint_weights(weights, params_dict) + + loader = AutoWeightsLoader( + self, + skip_prefixes=["lm_head."], + skip_substrs=["projector"], + ignore_unexpected_suffixes=[".bias", ".kv_scale"], + ) + return loader.load_weights(weights) + class Qwen3ForSequenceClassification(Qwen3ForPooledOutput): def __init__( diff --git a/test/manual/test_weight_loader_v2_equiv.py b/test/manual/test_weight_loader_v2_equiv.py index 1cf71c3e9c1a..051df1e47616 100644 --- a/test/manual/test_weight_loader_v2_equiv.py +++ b/test/manual/test_weight_loader_v2_equiv.py @@ -26,6 +26,7 @@ from sglang.srt.environ import envs MODEL = "Qwen/Qwen2-0.5B" +QWEN3_MODEL = "Qwen/Qwen3-0.6B" def _init_model_parallel() -> None: @@ -50,6 +51,10 @@ def _init_model_parallel() -> None: def _load_qwen2_native(v2: bool) -> torch.nn.Module: + return _load_native_model(MODEL, v2) + + +def _load_native_model(model_path: str, 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 @@ -58,7 +63,7 @@ def _load_qwen2_native(v2: bool) -> torch.nn.Module: from sglang.srt.utils import get_device server_args = ServerArgs( - model_path=MODEL, + model_path=model_path, dtype=torch.float16, trust_remote_code=True, ) @@ -106,6 +111,28 @@ def test_qwen2_v1_v2_state_dict_identical(self): msg=name, ) + @unittest.skipIf(not torch.cuda.is_available(), "needs GPU") + def test_qwen3_v1_v2_state_dict_identical(self): + model_v1 = _load_native_model(QWEN3_MODEL, v2=False) + state_v1 = _state_dict_cpu(model_v1) + del model_v1 + torch.cuda.empty_cache() + + model_v2 = _load_native_model(QWEN3_MODEL, 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 index 7ec68c9c3476..a43d6ef6c837 100644 --- a/test/registered/model_loading/test_weight_loader_v2_e2e.py +++ b/test/registered/model_loading/test_weight_loader_v2_e2e.py @@ -25,6 +25,7 @@ from sglang.test.test_utils import CustomTestCase MODEL = "Qwen/Qwen2-0.5B" +QWEN3_MODEL = "Qwen/Qwen3-0.6B" SHORT_PROMPT = "The capital of the United Kingdom is" @@ -65,6 +66,28 @@ def test_qwen2_native_v1_v2_generation_match(self): debug_text="qwen2 native v1 vs v2 weight loader", ) + def test_qwen3_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(QWEN3_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(QWEN3_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="qwen3 native v1 vs v2 weight loader", + ) + def test_transformers_impl_loads_and_generates(self): prompts = [SHORT_PROMPT] max_new_tokens = 16 diff --git a/test/registered/unit/model_loader/test_stacked_params_dispatch.py b/test/registered/unit/model_loader/test_stacked_params_dispatch.py new file mode 100644 index 000000000000..bd1d95848343 --- /dev/null +++ b/test/registered/unit/model_loader/test_stacked_params_dispatch.py @@ -0,0 +1,126 @@ +# 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_cpu_ci + +register_cpu_ci(est_time=6, suite="base-b-test-cpu") + +import unittest + +import torch +from torch import nn + +from sglang.srt.model_loader.auto_loader import ( + STANDARD_GATE_UP_MAPPING, + STANDARD_QKV_MAPPING, + StackedParamsDispatch, + load_with_stacked_dispatch, +) + + +class _ParamWithLoader(nn.Parameter): + def __new__(cls, data, weight_loader=None): + param = super().__new__(cls, data) + param.weight_loader = weight_loader or (lambda p, t, *a: p.copy_(t)) + return param + + +class TestStackedParamsDispatch(unittest.TestCase): + def test_try_load_qkv_shard(self): + qkv = _ParamWithLoader(torch.zeros(6, 4)) + params = {"qkv_proj.weight": qkv} + calls = [] + + def wl(param, tensor, shard_id): + calls.append(shard_id) + + qkv.weight_loader = wl + + target = STANDARD_QKV_MAPPING.try_load( + "q_proj.weight", torch.ones(2, 4), params + ) + self.assertEqual(target, "qkv_proj.weight") + self.assertEqual(calls, ["q"]) + + def test_try_load_returns_target_when_param_missing(self): + target = STANDARD_QKV_MAPPING.try_load("q_proj.weight", torch.ones(2, 4), {}) + self.assertEqual(target, "qkv_proj.weight") + + def test_try_load_no_match(self): + self.assertIsNone( + STANDARD_QKV_MAPPING.try_load("o_proj.weight", torch.ones(2, 4), {}) + ) + + def test_gate_up_mapping(self): + gate_up = _ParamWithLoader(torch.zeros(8, 4)) + params = {"gate_up_proj.weight": gate_up} + shard_ids = [] + + def wl(param, tensor, shard_id): + shard_ids.append(shard_id) + + gate_up.weight_loader = wl + + target = STANDARD_GATE_UP_MAPPING.try_load( + "up_proj.weight", torch.ones(4, 4), params + ) + self.assertEqual(target, "gate_up_proj.weight") + self.assertEqual(shard_ids, [1]) + + def test_load_with_stacked_dispatch_direct_param(self): + linear = nn.Linear(3, 2, bias=False) + linear.weight = nn.Parameter(torch.zeros(3, 2), requires_grad=False) + linear.weight.weight_loader = lambda p, t: p.data.copy_(t) + module = nn.Module() + module.down_proj = linear + + loaded = load_with_stacked_dispatch( + module, + [("down_proj.weight", torch.ones(3, 2))], + StackedParamsDispatch(mappings=()), + ) + self.assertIn("down_proj.weight", loaded) + self.assertTrue(torch.allclose(linear.weight.data, torch.ones(3, 2))) + + def test_load_with_stacked_dispatch_stacked_then_direct(self): + qkv_linear = nn.Linear(2, 6, bias=False) + o_linear = nn.Linear(2, 3, bias=False) + qkv_linear.weight = nn.Parameter(torch.zeros(6, 2), requires_grad=False) + o_linear.weight = nn.Parameter(torch.zeros(3, 2), requires_grad=False) + qkv_calls = [] + + def qkv_wl(param, tensor, shard_id): + qkv_calls.append(shard_id) + + qkv_linear.weight.weight_loader = qkv_wl + o_linear.weight.weight_loader = lambda p, t: p.data.copy_(t) + + module = nn.Module() + module.qkv_proj = qkv_linear + module.o_proj = o_linear + + loaded = load_with_stacked_dispatch( + module, + [ + ("q_proj.weight", torch.ones(2, 2)), + ("o_proj.weight", torch.full((3, 2), 2.0)), + ], + STANDARD_QKV_MAPPING, + ) + self.assertEqual(loaded, {"qkv_proj.weight", "o_proj.weight"}) + self.assertEqual(qkv_calls, ["q"]) + + +if __name__ == "__main__": + unittest.main() From 26e7454d074bb8c71ae46220d8d8980607544b6c Mon Sep 17 00:00:00 2001 From: JD-ETH Date: Wed, 15 Jul 2026 02:42:31 +0000 Subject: [PATCH 09/10] Add Weight Loader v2 standard MoE dispatch and migrations (PR3). ExpertParamsDispatch centralizes expert shard routing; gate standard MoE models behind SGLANG_ENABLE_WEIGHT_LOADER_V2 while keeping legacy loaders. Co-authored-by: Cursor --- python/sglang/srt/model_loader/auto_loader.py | 455 ++++++++++++++---- python/sglang/srt/models/mixtral.py | 54 +++ python/sglang/srt/models/olmoe.py | 21 + python/sglang/srt/models/phimoe.py | 28 ++ python/sglang/srt/models/qwen2_moe.py | 53 ++ python/sglang/srt/models/qwen3_moe.py | 48 ++ .../test_expert_params_dispatch.py | 91 ++++ 7 files changed, 663 insertions(+), 87 deletions(-) create mode 100644 test/registered/unit/model_loading/test_expert_params_dispatch.py diff --git a/python/sglang/srt/model_loader/auto_loader.py b/python/sglang/srt/model_loader/auto_loader.py index 81e7b9f0da38..f596673a11e2 100644 --- a/python/sglang/srt/model_loader/auto_loader.py +++ b/python/sglang/srt/model_loader/auto_loader.py @@ -11,21 +11,7 @@ # 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). -""" +"""Centralized weight loading utilities for native SGLang models.""" from __future__ import annotations @@ -45,44 +31,28 @@ "AutoWeightsLoader", "WeightsMapper", "StackedParamsDispatch", + "ExpertParamsDispatch", + "FusedExpertDispatch", "STANDARD_QKV_MAPPING", "STANDARD_GATE_UP_MAPPING", "STANDARD_STACKED_MAPPING", "LLAMA_STACKED_MAPPING", + "QWEN3_NEXT_GDN_STACKED_MAPPING", + "QWEN35_STACKED_MAPPING", + "QWEN35_MOE_IGNORE_SUFFIXES", + "MOE_EXPERT_STACKED_SKIP_SUBSTRS", + "try_load_stacked_skip_moe_experts", "load_with_stacked_dispatch", + "load_moe_sparse_block_weights", + "normalize_qwen35_weight_name", + "load_qwen35_moe_checkpoint_weights", "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( @@ -91,28 +61,18 @@ def try_load( 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"), @@ -138,7 +98,6 @@ def try_load( ) ) -# Llama-family full-path stacked mapping (dot-prefixed shard names). LLAMA_STACKED_MAPPING = StackedParamsDispatch( mappings=( (".qkv_proj", ".q_proj", "q"), @@ -149,6 +108,319 @@ def try_load( ) ) +QWEN3_NEXT_GDN_STACKED_MAPPING = StackedParamsDispatch( + mappings=( + ("in_proj_qkvz.", "in_proj_qkv.", (0, 1, 2)), + ("in_proj_qkvz.", "in_proj_z.", 3), + ("in_proj_ba.", "in_proj_b.", 0), + ("in_proj_ba.", "in_proj_a.", 1), + ) +) + +QWEN35_STACKED_MAPPING = StackedParamsDispatch( + mappings=STANDARD_STACKED_MAPPING.mappings + QWEN3_NEXT_GDN_STACKED_MAPPING.mappings +) + +QWEN35_MOE_IGNORE_SUFFIXES = ( + ".bias", + "_bias", + ".k_scale", + "_k_scale", + ".v_scale", + "_v_scale", + ".weight_scale", + "_weight_scale", + ".input_scale", + "_input_scale", +) + +MOE_EXPERT_STACKED_SKIP_SUBSTRS: tuple[str, ...] = ("mlp.experts", "experts.") + + +def try_load_stacked_skip_moe_experts( + dispatch: StackedParamsDispatch, + name: str, + tensor: torch.Tensor, + params_dict: dict[str, Parameter], + *, + skip_substrs: tuple[str, ...] = MOE_EXPERT_STACKED_SKIP_SUBSTRS, +) -> str | None: + for fused_name, source_name, shard_id in dispatch.mappings: + if source_name not in name: + continue + if any(skip in name for skip in skip_substrs): + continue + target = name.replace(source_name, fused_name) + param = params_dict.get(target) + if param is None: + return target + param.weight_loader(param, tensor, shard_id) + return target + return None + + +class ExpertParamsDispatch(msgspec.Struct, frozen=True): + mappings: tuple[tuple[str, str, int, str], ...] = () + + @classmethod + def from_fused_moe_mapping( + cls, + expert_params_mapping: list[tuple[str, str, int, str]], + ) -> ExpertParamsDispatch: + return cls(mappings=tuple(expert_params_mapping)) + + @classmethod + def from_gate_up_down( + cls, + *, + num_experts: int, + ckpt_gate_proj_name: str = "gate_proj", + ckpt_down_proj_name: str = "down_proj", + ckpt_up_proj_name: str = "up_proj", + ) -> ExpertParamsDispatch: + from sglang.srt.layers.moe.fused_moe_triton import FusedMoE + + return cls.from_fused_moe_mapping( + FusedMoE.make_expert_params_mapping( + ckpt_gate_proj_name=ckpt_gate_proj_name, + ckpt_down_proj_name=ckpt_down_proj_name, + ckpt_up_proj_name=ckpt_up_proj_name, + num_experts=num_experts, + ) + ) + + def try_load( + self, + name: str, + tensor: torch.Tensor, + params_dict: dict[str, Parameter], + ) -> str | None: + for param_name, weight_name, expert_id, shard_id in self.mappings: + if weight_name not in name: + continue + target = name.replace(weight_name, param_name) + if ( + target.endswith(".bias") or target.endswith("_bias") + ) and target not in params_dict: + return target + param = params_dict.get(target) + if param is None: + return target + weight_loader = getattr(param, "weight_loader", default_weight_loader) + try: + weight_loader( + param, + tensor, + target, + shard_id=shard_id, + expert_id=expert_id, + ) + except TypeError: + weight_loader(param, tensor) + return target + return None + + +class FusedExpertDispatch(msgspec.Struct, frozen=True): + num_experts: int + gate_up_ckpt_substr: str = "experts.gate_up_proj" + down_ckpt_substr: str = "experts.down_proj" + w13_runtime_substr: str = "experts.w13_weight" + w2_runtime_substr: str = "experts.w2_weight" + + @staticmethod + def fan_out_to_experts( + param: Parameter, + loaded_weight: torch.Tensor, + runtime_name: str, + shard_id: str, + num_experts: int, + ) -> bool: + weight_loader = getattr(param, "weight_loader", default_weight_loader) + for expert_id in range(num_experts): + weight_loader( + param, + loaded_weight[expert_id], + runtime_name, + shard_id, + expert_id, + ) + return True + + def try_load( + self, + name: str, + tensor: torch.Tensor, + params_dict: dict[str, Parameter], + ) -> str | None: + if self.gate_up_ckpt_substr in name: + target = name.replace(self.gate_up_ckpt_substr, self.w13_runtime_substr) + param = params_dict.get(target) + if param is None: + return target + w1, w3 = tensor.chunk(2, dim=-2) + self.fan_out_to_experts(param, w1, target, "w1", self.num_experts) + self.fan_out_to_experts(param, w3, target, "w3", self.num_experts) + return target + if self.down_ckpt_substr in name: + target = name.replace(self.down_ckpt_substr, self.w2_runtime_substr) + param = params_dict.get(target) + if param is None: + return target + self.fan_out_to_experts(param, tensor, target, "w2", self.num_experts) + return target + return None + + +def normalize_qwen35_weight_name(name: str) -> str: + if "language_model" in name: + name = name.replace(r"model.language_model.", r"model.") + if ".self_attn." in name: + name = name.replace(".self_attn", "") + return name + + +def load_qwen35_moe_checkpoint_weights( + module: nn.Module, + weights: Iterable[tuple[str, torch.Tensor]], + *, + num_experts: int, + expert_dispatch: ExpertParamsDispatch, + fused_dispatch: FusedExpertDispatch | None, + stacked_mapping: StackedParamsDispatch = QWEN35_STACKED_MAPPING, + ignore_suffixes: tuple[str, ...] = QWEN35_MOE_IGNORE_SUFFIXES, + start_layer: int | None = None, + end_layer: int | None = None, + skip_substrs: tuple[str, ...] = ("rotary_emb.inv_freq", "mtp", "visual"), + enable_shared_expert_fusion: bool = False, + shared_expert_slot: int | None = None, + encoder_only: bool = False, + remap_visual: bool = False, + on_embed_for_tied_lm_head: Callable[[str, torch.Tensor], None] | None = None, +) -> set[str]: + import logging + + log = logging.getLogger(__name__) + loaded_params: set[str] = set() + params_dict = dict(module.named_parameters(remove_duplicate=False)) + is_fused_expert = False + active_expert_dispatch = expert_dispatch + fused_gate_up_mapping: tuple[tuple[str, str, int, str], ...] = ( + ("experts.w13_weight", "experts.gate_up_proj", 0, "w1"), + ("experts.w2_weight", "experts.down_proj", 0, "w2"), + ) + if enable_shared_expert_fusion and shared_expert_slot is not None: + slot = shared_expert_slot + fused_gate_up_mapping = fused_gate_up_mapping + ( + ("experts.w13_", f"experts.{slot}.gate_up_proj.", slot, "w1"), + ("experts.w2_", f"experts.{slot}.down_proj.", slot, "w2"), + ("experts.w13_", f"experts.{slot}.gate_proj.", slot, "w1"), + ("experts.w13_", f"experts.{slot}.up_proj.", slot, "w3"), + ) + for name, loaded_weight in weights: + if any(sub in name for sub in skip_substrs): + continue + name = normalize_qwen35_weight_name(name) + if on_embed_for_tied_lm_head is not None: + on_embed_for_tied_lm_head(name, loaded_weight) + if enable_shared_expert_fusion and shared_expert_slot is not None: + if "mlp.shared_expert." in name: + name = name.replace( + "mlp.shared_expert.", + f"mlp.experts.{shared_expert_slot}.", + ) + layer_id = get_layer_id(name) + if ( + layer_id is not None + and start_layer is not None + and end_layer is not None + and (layer_id < start_layer or layer_id >= end_layer) + ): + continue + if ( + "experts.gate_up_proj" in name + or "experts.down_proj" in name + or name.endswith("experts.gate_up_proj") + or name.endswith("experts.down_proj") + ): + is_fused_expert = True + active_expert_dispatch = ExpertParamsDispatch.from_fused_moe_mapping( + list(fused_gate_up_mapping) + ) + target = try_load_stacked_skip_moe_experts( + stacked_mapping, name, loaded_weight, params_dict + ) + if target is not None: + loaded_params.add(name) + continue + if is_fused_expert and fused_dispatch is not None: + fused_target = fused_dispatch.try_load(name, loaded_weight, params_dict) + if fused_target is not None: + loaded_params.add(name) + continue + if enable_shared_expert_fusion and shared_expert_slot is not None: + handled = False + for ( + param_name, + weight_name, + expert_id, + shard_id, + ) in fused_gate_up_mapping: + if weight_name not in name: + continue + if "visual" in name or encoder_only: + continue + name_mapped = name.replace(weight_name, param_name) + if name_mapped not in params_dict: + handled = True + break + param = params_dict[name_mapped] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + if f"{shared_expert_slot}.gate_up_proj" in name: + w1, w3 = loaded_weight.chunk(2, dim=-2) + weight_loader(param, w1, name_mapped, "w1", expert_id) + weight_loader(param, w3, name_mapped, "w3", expert_id) + else: + weight_loader( + param, + loaded_weight, + name_mapped, + shard_id, + expert_id, + ) + loaded_params.add(name) + handled = True + break + if handled: + continue + expert_target = active_expert_dispatch.try_load( + name, loaded_weight, params_dict + ) + if expert_target is not None: + loaded_params.add(name) + continue + if any( + weight_name in name + for _, weight_name, _, _ in active_expert_dispatch.mappings + ): + loaded_params.add(name) + continue + if remap_visual and "visual" in name: + name = name.replace(r"attn.qkv.", r"attn.qkv_proj.") + name = name.replace(r"model.visual.", r"visual.") + if name.endswith(ignore_suffixes) and name not in params_dict: + continue + if name in params_dict: + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + else: + log.warning("Parameter %s not found in params_dict", name) + return loaded_params + def load_with_stacked_dispatch( module: nn.Module, @@ -163,7 +435,14 @@ def load_with_stacked_dispatch( for name, tensor in weights: target = mapping.try_load(name, tensor, params_dict) if target is not None: - loaded.add(target) + if target in params_dict: + loaded.add(target) + continue + if name.endswith("_scale") and name not in params_dict: + if abs(tensor.item() - 1.0) >= 1e-6: + raise AssertionError( + f"Expected unit scale 1.0, got {tensor.item()} for {name}" + ) continue if name in params_dict: wl = getattr(params_dict[name], "weight_loader", default_weight_loader) @@ -174,9 +453,42 @@ def load_with_stacked_dispatch( return loaded -# --------------------------------------------------------------------------- -# Pipeline Parallel Weight Filter -# --------------------------------------------------------------------------- +def load_moe_sparse_block_weights( + module: nn.Module, + weights: Iterable[tuple[str, torch.Tensor]], + *, + expert_dispatch: ExpertParamsDispatch, + dense_stacked: StackedParamsDispatch = STANDARD_GATE_UP_MAPPING, +) -> set[str]: + loaded: set[str] = set() + params_dict = dict(module.named_parameters()) + for name, tensor in weights: + target = try_load_stacked_skip_moe_experts( + dense_stacked, name, tensor, params_dict + ) + if target is not None: + if target in params_dict: + loaded.add(target) + continue + target = expert_dispatch.try_load(name, tensor, params_dict) + if target is not None: + if target in params_dict: + loaded.add(target) + continue + if name.endswith("_scale") and name not in params_dict: + if abs(tensor.item() - 1.0) >= 1e-6: + raise AssertionError( + f"Expected unit scale 1.0, got {tensor.item()} for {name}" + ) + continue + if name.endswith(".bias") and name not in params_dict: + continue + if name not in params_dict: + continue + wl = getattr(params_dict[name], "weight_loader", default_weight_loader) + wl(params_dict[name], tensor) + loaded.add(name) + return loaded def filter_pp_weights( @@ -184,11 +496,6 @@ def filter_pp_weights( 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): @@ -196,29 +503,10 @@ def filter_pp_weights( 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 @@ -228,21 +516,14 @@ def decorator(fn: Callable[[nn.Module], WeightsMapper]): 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", diff --git a/python/sglang/srt/models/mixtral.py b/python/sglang/srt/models/mixtral.py index db258623fd3a..79a68af22e55 100644 --- a/python/sglang/srt/models/mixtral.py +++ b/python/sglang/srt/models/mixtral.py @@ -78,6 +78,7 @@ def __init__( super().__init__() self.tp_size = get_tensor_model_parallel_world_size() self.hidden_size = hidden_size + self._ckpt_num_experts = num_experts # Gate always runs at half / full precision for now. self.gate = ReplicatedLinear( @@ -117,6 +118,26 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) return final_hidden_states.view(orig_shape) + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> set[str]: + from sglang.srt.model_loader.auto_loader import ( + ExpertParamsDispatch, + StackedParamsDispatch, + load_moe_sparse_block_weights, + ) + + expert = ExpertParamsDispatch.from_gate_up_down( + num_experts=self._ckpt_num_experts, + ckpt_gate_proj_name="w1", + ckpt_down_proj_name="w2", + ckpt_up_proj_name="w3", + ) + return load_moe_sparse_block_weights( + self, + weights, + expert_dispatch=expert, + dense_stacked=StackedParamsDispatch(), + ) + class MixtralAttention(nn.Module): def __init__( @@ -198,6 +219,14 @@ def forward( 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, + load_with_stacked_dispatch, + ) + + return load_with_stacked_dispatch(self, weights, STANDARD_QKV_MAPPING) + class MixtralDecoderLayer(nn.Module): def __init__( @@ -388,6 +417,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"), @@ -476,5 +512,23 @@ 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]: + from sglang.srt.model_loader.auto_loader import ( + AutoWeightsLoader, + filter_pp_weights, + ) + + if hasattr(self.model, "start_layer"): + weights = filter_pp_weights( + weights, self.model.start_layer, self.model.end_layer + ) + + loader = AutoWeightsLoader( + self, + skip_substrs=["rotary_emb.inv_freq"], + ignore_unexpected_suffixes=[".bias", ".kv_scale"], + ) + return loader.load_weights(weights) + EntryClass = MixtralForCausalLM diff --git a/python/sglang/srt/models/olmoe.py b/python/sglang/srt/models/olmoe.py index 46be4fe88edf..b26c12fed1c1 100644 --- a/python/sglang/srt/models/olmoe.py +++ b/python/sglang/srt/models/olmoe.py @@ -347,6 +347,13 @@ def forward( ) 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"), @@ -435,5 +442,19 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): ) weight_loader(param, loaded_weight) + def _load_weights_v2(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> set[str]: + from sglang.srt.model_loader.auto_loader import ( + filter_pp_weights, + load_standard_moe_root_weights_v2, + ) + + if hasattr(self.model, "start_layer"): + weights = filter_pp_weights( + weights, self.model.start_layer, self.model.end_layer + ) + return load_standard_moe_root_weights_v2( + self, weights, num_experts=self.config.num_experts + ) + EntryClass = OlmoeForCausalLM diff --git a/python/sglang/srt/models/phimoe.py b/python/sglang/srt/models/phimoe.py index a359483de3ef..ba4b99d3e6e8 100644 --- a/python/sglang/srt/models/phimoe.py +++ b/python/sglang/srt/models/phimoe.py @@ -499,6 +499,13 @@ def forward( return self.pooler(hidden_states, forward_batch) 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"), @@ -555,5 +562,26 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): ) weight_loader(param, loaded_weight) + def _load_weights_v2(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> set[str]: + from sglang.srt.model_loader.auto_loader import ( + STANDARD_QKV_MAPPING, + filter_pp_weights, + load_standard_moe_root_weights_v2, + ) + + if hasattr(self.model, "start_layer"): + weights = filter_pp_weights( + weights, self.model.start_layer, self.model.end_layer + ) + return load_standard_moe_root_weights_v2( + self, + weights, + num_experts=self.config.num_local_experts, + ckpt_gate_proj_name="w1", + ckpt_down_proj_name="w2", + ckpt_up_proj_name="w3", + stacked_mapping=STANDARD_QKV_MAPPING, + ) + EntryClass = PhiMoEForCausalLM diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index 324e8678d8c7..6c25b6864e02 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -223,6 +223,14 @@ def forward( ) 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, + load_with_stacked_dispatch, + ) + + return load_with_stacked_dispatch(self, weights, STANDARD_GATE_UP_MAPPING) + class Qwen2MoeSparseMoeBlock(nn.Module): def __init__( @@ -246,6 +254,7 @@ def __init__( f"the number of experts {config.num_experts}." ) self.num_experts = config.num_experts + self._ckpt_num_experts = config.num_experts self.num_shared_experts = get_num_shared_experts(config) self.num_fused_shared_experts = 0 @@ -592,6 +601,17 @@ def forward( return final_hidden_states.view(num_tokens, hidden_dim) + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> set[str]: + from sglang.srt.model_loader.auto_loader import ( + ExpertParamsDispatch, + load_moe_sparse_block_weights, + ) + + expert = ExpertParamsDispatch.from_gate_up_down( + num_experts=self._ckpt_num_experts, + ) + return load_moe_sparse_block_weights(self, weights, expert_dispatch=expert) + class Qwen2MoeAttention(nn.Module): def __init__( @@ -688,6 +708,14 @@ def forward( 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, + load_with_stacked_dispatch, + ) + + return load_with_stacked_dispatch(self, weights, STANDARD_QKV_MAPPING) + class Qwen2MoeDecoderLayer(nn.Module): def __init__( @@ -1077,6 +1105,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"), @@ -1162,6 +1197,24 @@ 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]: + from sglang.srt.model_loader.auto_loader import ( + AutoWeightsLoader, + filter_pp_weights, + ) + + if hasattr(self.model, "start_layer"): + weights = filter_pp_weights( + weights, self.model.start_layer, self.model.end_layer + ) + + loader = AutoWeightsLoader( + self, + skip_substrs=["rotary_emb.inv_freq"], + ignore_unexpected_suffixes=[".bias"], + ) + return loader.load_weights(weights) + @classmethod def get_model_config_for_expert_location(cls, config): return ModelConfigForExpertLocation( diff --git a/python/sglang/srt/models/qwen3_moe.py b/python/sglang/srt/models/qwen3_moe.py index 5c6c76bf13fa..a7ef7da17273 100644 --- a/python/sglang/srt/models/qwen3_moe.py +++ b/python/sglang/srt/models/qwen3_moe.py @@ -248,6 +248,8 @@ def __init__( f"the number of experts {config.num_experts}." ) + self._ckpt_num_experts = config.num_experts + from sglang.srt.layers.quantization.gguf import GGUFConfig norm_topk_prob = getattr(config, "norm_topk_prob", True) @@ -437,6 +439,17 @@ def op_combine_b(self, state): def op_output(self, state): state.hidden_states_mlp_output = state.pop("hidden_states_after_combine") + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> set[str]: + from sglang.srt.model_loader.auto_loader import ( + ExpertParamsDispatch, + load_moe_sparse_block_weights, + ) + + expert = ExpertParamsDispatch.from_gate_up_down( + num_experts=self._ckpt_num_experts, + ) + return load_moe_sparse_block_weights(self, weights, expert_dispatch=expert) + class Qwen3MoeAttention(nn.Module): def __init__( @@ -717,6 +730,14 @@ def forward( ) return self.forward_core(s) + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> set[str]: + from sglang.srt.model_loader.auto_loader import ( + STANDARD_QKV_MAPPING, + load_with_stacked_dispatch, + ) + + return load_with_stacked_dispatch(self, weights, STANDARD_QKV_MAPPING) + class Qwen3MoeDecoderLayer(nn.Module): def __init__( @@ -1109,6 +1130,15 @@ def set_dflash_layers_to_capture(self, layer_ids: List[int]): def load_weights( self, weights: Iterable[Tuple[str, torch.Tensor]], is_mtp: bool = False + ): + from sglang.srt.environ import envs + + if envs.SGLANG_ENABLE_WEIGHT_LOADER_V2.get() and not is_mtp: + return self._load_weights_v2(weights) + return self._legacy_load_weights(weights, is_mtp=is_mtp) + + def _legacy_load_weights( + self, weights: Iterable[Tuple[str, torch.Tensor]], is_mtp: bool = False ): stacked_params_mapping = [ # (param_name, shard_name, shard_id) @@ -1239,6 +1269,24 @@ def load_weights( } ) + def _load_weights_v2(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> set[str]: + from sglang.srt.model_loader.auto_loader import ( + AutoWeightsLoader, + filter_pp_weights, + ) + + if hasattr(self.model, "start_layer"): + weights = filter_pp_weights( + weights, self.model.start_layer, self.model.end_layer + ) + + loader = AutoWeightsLoader( + self, + skip_substrs=["rotary_emb.inv_freq", "mtp"], + ignore_unexpected_suffixes=[".bias"], + ) + return loader.load_weights(weights) + @classmethod def get_model_config_for_expert_location(cls, config): return ModelConfigForExpertLocation( diff --git a/test/registered/unit/model_loading/test_expert_params_dispatch.py b/test/registered/unit/model_loading/test_expert_params_dispatch.py new file mode 100644 index 000000000000..12b056a48eda --- /dev/null +++ b/test/registered/unit/model_loading/test_expert_params_dispatch.py @@ -0,0 +1,91 @@ +import unittest +from unittest import mock + +import torch + +from sglang.srt.layers.moe.fused_moe_triton import FusedMoE +from sglang.srt.model_loader.auto_loader import ( + MOE_EXPERT_STACKED_SKIP_SUBSTRS, + STANDARD_GATE_UP_MAPPING, + ExpertParamsDispatch, + try_load_stacked_skip_moe_experts, +) +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=8, suite="base-b-test-cpu") + + +class _RecordingLoader: + def __init__(self): + self.calls: list[tuple] = [] + + def __call__(self, param, tensor, name, *, shard_id, expert_id): + self.calls.append((name, shard_id, expert_id, tuple(tensor.shape))) + + +class TestExpertParamsDispatch(unittest.TestCase): + def test_expert_routing_gate_proj(self): + mapping = FusedMoE.make_expert_params_mapping( + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=2, + ) + dispatch = ExpertParamsDispatch.from_fused_moe_mapping(mapping) + loader = _RecordingLoader() + param = mock.Mock() + param.weight_loader = loader + params = {"model.layers.0.mlp.experts.w13_weight": param} + tensor = torch.zeros(4, 4) + ckpt_name = "model.layers.0.mlp.experts.1.gate_proj.weight" + target = dispatch.try_load(ckpt_name, tensor, params) + self.assertEqual(target, "model.layers.0.mlp.experts.w13_weight") + self.assertEqual(len(loader.calls), 1) + self.assertEqual(loader.calls[0][1], "w1") + self.assertEqual(loader.calls[0][2], 1) + + def test_from_gate_up_down_mixtral_names(self): + dispatch = ExpertParamsDispatch.from_gate_up_down( + num_experts=1, + ckpt_gate_proj_name="w1", + ckpt_down_proj_name="w2", + ckpt_up_proj_name="w3", + ) + self.assertTrue( + any( + "experts.0.w1." in weight_name + for _, weight_name, _, _ in dispatch.mappings + ) + ) + + def test_stacked_skip_moe_experts_before_rename(self): + gate_up = mock.Mock() + gate_up.weight_loader = mock.Mock() + params = {"mlp.experts.0.gate_up_proj.weight": gate_up} + tensor = torch.zeros(2, 2) + name = "mlp.experts.0.gate_proj.weight" + result = try_load_stacked_skip_moe_experts( + STANDARD_GATE_UP_MAPPING, name, tensor, params + ) + self.assertIsNone(result) + gate_up.weight_loader.assert_not_called() + + def test_stacked_applies_to_shared_expert_path(self): + gate_up = mock.Mock() + gate_up.weight_loader = mock.Mock() + params = {"mlp.shared_expert.gate_up_proj.weight": gate_up} + tensor = torch.zeros(2, 2) + name = "mlp.shared_expert.gate_proj.weight" + result = try_load_stacked_skip_moe_experts( + STANDARD_GATE_UP_MAPPING, name, tensor, params + ) + self.assertEqual(result, "mlp.shared_expert.gate_up_proj.weight") + gate_up.weight_loader.assert_called_once() + + def test_moe_expert_skip_substrs_cover_experts_prefix(self): + self.assertIn("mlp.experts", MOE_EXPERT_STACKED_SKIP_SUBSTRS) + self.assertIn("experts.", MOE_EXPERT_STACKED_SKIP_SUBSTRS) + + +if __name__ == "__main__": + unittest.main() From d5f3349cf57fc479648688e2407d0dabe0104624 Mon Sep 17 00:00:00 2001 From: JD-ETH Date: Wed, 15 Jul 2026 02:52:44 +0000 Subject: [PATCH 10/10] Add weight loader v2 for DeepSeek MLA mixin (PR4). Opt-in via SGLANG_ENABLE_WEIGHT_LOADER_V2 with AutoWeightsLoader submodule loaders, shared-expert preload, MLA fusion remaps, and post_load_weights for w_kc/w_vc. Co-authored-by: Cursor --- python/sglang/srt/model_loader/auto_loader.py | 113 +++++++++++++- .../deepseek_common/deepseek_weight_loader.py | 92 ++++++++++- python/sglang/srt/models/deepseek_v2.py | 86 +++++++++++ python/sglang/srt/models/glm4_moe_lite.py | 4 + .../test_weight_loader_v2_deepseek.py | 144 ++++++++++++++++++ 5 files changed, 433 insertions(+), 6 deletions(-) create mode 100644 test/registered/unit/model_loader/test_weight_loader_v2_deepseek.py diff --git a/python/sglang/srt/model_loader/auto_loader.py b/python/sglang/srt/model_loader/auto_loader.py index f596673a11e2..763b81288784 100644 --- a/python/sglang/srt/model_loader/auto_loader.py +++ b/python/sglang/srt/model_loader/auto_loader.py @@ -15,6 +15,8 @@ from __future__ import annotations +import re + from collections.abc import Callable, Iterable from typing import Union @@ -49,7 +51,11 @@ "filter_pp_weights", "register_weight_remap", "get_weight_remap", -] + + "MultiInputFusion", + "DEEPSEEK_GATE_UP_MAPPING", + "remap_fused_shared_expert_names", + "maybe_remap_deepseek_mla_kv_scale",] class StackedParamsDispatch(msgspec.Struct, frozen=True): @@ -530,3 +536,108 @@ def _llama_remap(model: nn.Module) -> WeightsMapper: ".weight_scale_inv": ".weight_scale", } ) + +DEEPSEEK_GATE_UP_MAPPING = STANDARD_GATE_UP_MAPPING + + +_SHARED_EXPERT_PATTERN = re.compile( + r"^model\.layers\.(\d+)\.mlp\.shared_experts\.(.+)$" +) + + +class MultiInputFusion: + def __init__( + self, + *, + source_substrs: tuple[str, ...], + fused_param_substr: str, + cat_dim: int = 0, + tensor_cloner: Callable[[torch.Tensor], torch.Tensor] | None = None, + ) -> None: + self.source_substrs = source_substrs + self.fused_param_substr = fused_param_substr + self.cat_dim = cat_dim + self.tensor_cloner = tensor_cloner + self._pending: dict[str, dict[str, torch.Tensor]] = {} + + def try_load( + self, + name: str, + tensor: torch.Tensor, + params_dict: dict[str, Parameter], + ) -> str | None: + matched = next((s for s in self.source_substrs if s in name), None) + if matched is None: + return None + if self.tensor_cloner is not None: + tensor = self.tensor_cloner(tensor) + group_key = name.split(matched, 1)[0] + bucket = self._pending.setdefault(group_key, {}) + bucket[matched] = tensor + if len(bucket) < len(self.source_substrs): + return None + parts = [bucket[s] for s in self.source_substrs] + fused = ( + parts[0] + if all(p.shape == torch.Size([]) for p in parts) + else torch.cat(parts, dim=self.cat_dim) + ) + fused_name = ( + name.replace("q_a_proj", self.fused_param_substr) + if matched == "q_a_proj" + else name.replace("kv_a_proj_with_mqa", self.fused_param_substr) + ) + param = params_dict.get(fused_name) + if param is None: + self._pending.pop(group_key, None) + return fused_name + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, fused) + self._pending.pop(group_key, None) + return fused_name + + +def remap_fused_shared_expert_names( + weights: Iterable[tuple[str, torch.Tensor]], + n_routed_experts: int, +) -> Iterable[tuple[str, torch.Tensor]]: + for name, tensor in weights: + match = _SHARED_EXPERT_PATTERN.match(name) + if match: + layer_id, suffix = match.group(1), match.group(2) + name = f"model.layers.{layer_id}.mlp.experts.{n_routed_experts}.{suffix}" + elif "mlp.shared_experts" in name: + name = name.replace("mlp.shared_experts", f"mlp.experts.{n_routed_experts}") + yield name, tensor + + +def maybe_remap_deepseek_mla_kv_scale( + name: str, + params_dict: dict[str, Parameter], +) -> str | None: + if name in params_dict: + return name + if "k_scale" not in name and "v_scale" not in name: + return name + for scale in ("k_scale", "v_scale"): + if scale in name: + if "kv_a_proj_with_mqa" in name: + candidate = name.replace("kv_a_proj_with_mqa", "attn_mqa") + else: + candidate = name.replace(f"{scale[0]}_proj", "attn_mqa") + if candidate in params_dict: + return candidate + return name + + +@register_weight_remap( + "DeepseekV2ForCausalLM", + "DeepseekV3ForCausalLM", + "Glm4MoeLiteForCausalLM", +) +def _deepseek_mla_remap(model: nn.Module) -> WeightsMapper: + n_routed = getattr(model.config, "n_routed_experts", None) + substr_map: dict[str, str | None] = {} + if n_routed is not None: + substr_map["mlp.shared_experts"] = f"mlp.experts.{n_routed}" + return WeightsMapper(orig_to_new_substr=substr_map) diff --git a/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py b/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py index 754dc1bbacd0..7b6292b42cfc 100644 --- a/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py +++ b/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py @@ -107,12 +107,15 @@ def do_load_weights( weights: Iterable[Tuple[str, torch.Tensor]], is_nextn: bool = False, ): - """Load model weights from checkpoint. + if envs.SGLANG_ENABLE_WEIGHT_LOADER_V2.get(): + return self._do_load_weights_v2(weights, is_nextn=is_nextn) + return self._legacy_do_load_weights(weights, is_nextn=is_nextn) - Args: - weights: Iterable of (weight_name, weight_tensor) pairs - is_nextn: Whether loading NextN speculative decoding weights - """ + def _legacy_do_load_weights( + self, + weights: Iterable[Tuple[str, torch.Tensor]], + is_nextn: bool = False, + ): nextn_conf = self._initialize_nextn_conf(is_nextn) weights = self._maybe_quant_weights_to_fp8_ue8m0( @@ -374,6 +377,85 @@ def do_load_weights( self.post_load_weights(is_nextn=is_nextn, weight_names=weight_names) + def _do_load_weights_v2( + self, + weights: Iterable[Tuple[str, torch.Tensor]], + is_nextn: bool = False, + ) -> set[str]: + from sglang.srt.model_loader.auto_loader import ( + AutoWeightsLoader, + filter_pp_weights, + get_weight_remap, + remap_fused_shared_expert_names, + ) + + nextn_conf = self._initialize_nextn_conf(is_nextn) + weights = self._maybe_quant_weights_to_fp8_ue8m0( + weights, NVFP4_CKPT_FP8_ATTN_QUANT_MODULES, nextn_conf + ) + if self.num_fused_shared_experts > 0: + assert self.num_fused_shared_experts == 1 + log_info_on_rank0(logger, "Shared experts fusion optimization enabled.") + weights = remap_fused_shared_expert_names( + weights, self.config.n_routed_experts + ) + weight_names: list[str] = [] + + def _preprocess() -> Iterable[Tuple[str, torch.Tensor]]: + for name, loaded_weight in weights: + layer_id = get_layer_id(name) + if ( + layer_id is not None + and hasattr(self.model, "start_layer") + and ( + layer_id < self.model.start_layer + or layer_id >= self.model.end_layer + ) + ): + continue + match nextn_conf: + case NextNEnabledConfig( + nextn_layer_prefix=layer_prefix, + nextn_spec_weight_names=spec_weight_names, + ): + if not name.startswith(layer_prefix): + continue + if "shared_head.head" in name or "embed_tokens" in name: + continue + if any(s in name for s in spec_weight_names): + name = name.replace(layer_prefix, "model") + else: + name = name.replace(layer_prefix, "model.decoder") + case NextNDisabledConfig(): + if hasattr(self.config, "num_nextn_predict_layers"): + num_nextn_layers = self.config.num_nextn_predict_layers + if num_nextn_layers > 0 and name.startswith("model.layers"): + name_list = name.split(".") + if ( + len(name_list) >= 3 + and int(name_list[2]) + >= self.config.num_hidden_layers + ): + continue + if "rotary_emb.inv_freq" in name: + continue + if ".embed_tokens." in name and not self.pp_group.is_first_rank: + continue + if ".norm." in name and not self.pp_group.is_last_rank: + continue + weight_names.append(name) + yield name, loaded_weight + + preprocessed = _preprocess() + if hasattr(self.model, "start_layer"): + preprocessed = filter_pp_weights( + preprocessed, self.model.start_layer, self.model.end_layer + ) + loader = AutoWeightsLoader(self, ignore_unexpected_suffixes=[".bias"]) + loaded = loader.load_weights(preprocessed, mapper=get_weight_remap(self)) + self.post_load_weights(is_nextn=is_nextn, weight_names=weight_names) + return loaded + def _initialize_nextn_conf(self, is_nextn: bool) -> NextNConfig: """ Initialize the nextn configuration. diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index c6208ff1d1e6..59409ac3fe60 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -270,6 +270,14 @@ def __init__( self._fused_clamp_fp8_checked = False self._fused_clamp_use_fp8 = False + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> set[str]: + from sglang.srt.model_loader.auto_loader import ( + STANDARD_GATE_UP_MAPPING, + load_with_stacked_dispatch, + ) + + return load_with_stacked_dispatch(self, weights, STANDARD_GATE_UP_MAPPING) + def forward( self, x, @@ -817,6 +825,36 @@ def get_moe_weights(self): ) ] + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> set[str]: + from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE + from sglang.srt.model_loader.auto_loader import ( + DEEPSEEK_GATE_UP_MAPPING, + ExpertParamsDispatch, + load_moe_sparse_block_weights, + ) + from sglang.srt.models.deepseek_common.utils import _is_npu + + expert_mapping = FusedMoE.make_expert_params_mapping( + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.n_routed_experts + self.num_fused_shared_experts, + ) + quant_config = getattr(self.experts, "quant_config", None) + if quant_config is not None and quant_config.get_name() == "w4afp8": + expert_mapping += FusedMoE.make_expert_input_scale_params_mapping( + num_experts=self.config.n_routed_experts + ) + expert_dispatch = ExpertParamsDispatch.from_fused_moe_mapping(expert_mapping) + if _is_npu: + weights = ((n.replace("weight_packed", "weight"), w) for n, w in weights) + return load_moe_sparse_block_weights( + self, + weights, + expert_dispatch=expert_dispatch, + dense_stacked=DEEPSEEK_GATE_UP_MAPPING, + ) + def forward( self, hidden_states: torch.Tensor, @@ -1775,6 +1813,54 @@ def op_core(self, state): else: state.hidden_states_after_attn = result + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> set[str]: + from sglang.srt.model_loader.auto_loader import ( + MultiInputFusion, + maybe_remap_deepseek_mla_kv_scale, + ) + from sglang.srt.model_loader.weight_utils import default_weight_loader + from sglang.srt.models.deepseek_common.deepseek_weight_loader import ( + _clone_if_runai_streamed_tensor, + ) + from sglang.srt.models.deepseek_common.utils import _is_npu + + loaded: set[str] = set() + params_dict = dict(self.named_parameters()) + mla_fusion = None + if self.q_lora_rank is not None: + cat_dim = 0 + if self.quant_config is not None and self.quant_config.get_name() in ( + "awq", + "awq_marlin", + "moe_wna16", + ): + cat_dim = 1 + mla_fusion = MultiInputFusion( + source_substrs=("q_a_proj", "kv_a_proj_with_mqa"), + fused_param_substr="fused_qkv_a_proj_with_mqa", + cat_dim=cat_dim, + tensor_cloner=_clone_if_runai_streamed_tensor, + ) + for name, tensor in weights: + if _is_npu: + name = name.replace("weight_packed", "weight") + if mla_fusion is not None: + fused_target = mla_fusion.try_load(name, tensor, params_dict) + if fused_target is not None: + loaded.add(fused_target) + continue + if "q_a_proj" in name or "kv_a_proj_with_mqa" in name: + continue + name = maybe_remap_deepseek_mla_kv_scale(name, params_dict) + if name is None or (name.endswith(".bias") and name not in params_dict): + continue + if name not in params_dict: + continue + wl = getattr(params_dict[name], "weight_loader", default_weight_loader) + wl(params_dict[name], tensor) + loaded.add(name) + return loaded + def forward( self, positions: torch.Tensor, diff --git a/python/sglang/srt/models/glm4_moe_lite.py b/python/sglang/srt/models/glm4_moe_lite.py index a36d4ee86fc6..37a4d60d20e6 100644 --- a/python/sglang/srt/models/glm4_moe_lite.py +++ b/python/sglang/srt/models/glm4_moe_lite.py @@ -35,6 +35,7 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import ( use_symmetric_memory, ) +from sglang.srt.environ import envs from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo @@ -1054,6 +1055,9 @@ def load_weights( is_nextn=False, params_dict=None, ): + if envs.SGLANG_ENABLE_WEIGHT_LOADER_V2.get(): + return self.do_load_weights(weights, is_nextn=is_nextn) + if is_nextn: if hasattr(self.config, "num_nextn_predict_layers"): num_nextn_layers = self.config.num_nextn_predict_layers diff --git a/test/registered/unit/model_loader/test_weight_loader_v2_deepseek.py b/test/registered/unit/model_loader/test_weight_loader_v2_deepseek.py new file mode 100644 index 000000000000..fc5525068a35 --- /dev/null +++ b/test/registered/unit/model_loader/test_weight_loader_v2_deepseek.py @@ -0,0 +1,144 @@ +# Copyright 2026 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. +# ============================================================================== + +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch +from torch.nn import Parameter + +from sglang.srt.model_loader.auto_loader import ( + ExpertParamsDispatch, + MultiInputFusion, + _deepseek_mla_remap, + maybe_remap_deepseek_mla_kv_scale, + remap_fused_shared_expert_names, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=12, suite="base-a-test-cpu") + + +class TestDeepseekSharedExpertRemap(CustomTestCase): + def test_remap_fused_shared_expert_names(self): + weights = [ + ( + "model.layers.3.mlp.shared_experts.gate_proj.weight", + torch.zeros(2), + ), + ] + out = list(remap_fused_shared_expert_names(weights, n_routed_experts=256)) + self.assertEqual( + out[0][0], + "model.layers.3.mlp.experts.256.gate_proj.weight", + ) + + def test_registry_shared_expert_substr(self): + model = SimpleNamespace(config=SimpleNamespace(n_routed_experts=128)) + mapper = _deepseek_mla_remap(model) + mapped = list( + mapper.apply( + [ + ( + "model.layers.0.mlp.shared_experts.down_proj.weight", + torch.zeros(1), + ) + ] + ) + ) + self.assertEqual( + mapped[0][0], + "model.layers.0.mlp.experts.128.down_proj.weight", + ) + + +class TestDeepseekMlaKvScaleRemap(CustomTestCase): + def test_maybe_remap_deepseek_mla_kv_scale(self): + params = { + "model.layers.0.self_attn.attn_mqa.k_scale": Parameter(torch.ones(1)), + } + name = "model.layers.0.self_attn.kv_a_proj_with_mqa.k_scale" + self.assertEqual( + maybe_remap_deepseek_mla_kv_scale(name, params), + "model.layers.0.self_attn.attn_mqa.k_scale", + ) + + +class TestMultiInputFusionMla(CustomTestCase): + def test_fuses_q_and_kv_a_proj(self): + loaded: list[torch.Tensor] = [] + + def _wl(param, weight): + loaded.append(weight.clone()) + + fused = Parameter(torch.zeros(6, 4)) + fused.weight_loader = _wl + prefix = "model.layers.0.self_attn" + params = {f"{prefix}.fused_qkv_a_proj_with_mqa.weight": fused} + fusion = MultiInputFusion( + source_substrs=("q_a_proj", "kv_a_proj_with_mqa"), + fused_param_substr="fused_qkv_a_proj_with_mqa", + cat_dim=0, + ) + q = torch.arange(8, dtype=torch.float32).reshape(4, 2) + kv = torch.arange(8, 16, dtype=torch.float32).reshape(4, 2) + self.assertIsNone(fusion.try_load(f"{prefix}.q_a_proj.weight", q, params)) + r2 = fusion.try_load(f"{prefix}.kv_a_proj_with_mqa.weight", kv, params) + self.assertEqual(r2, f"{prefix}.fused_qkv_a_proj_with_mqa.weight") + torch.testing.assert_close(loaded[0], torch.cat([q, kv], dim=0)) + + +class TestExpertParamsDispatch(CustomTestCase): + def test_routes_expert_gate_proj(self): + calls = [] + + def _wl(param, tensor, qualname, shard_id=None, expert_id=None): + calls.append((qualname, shard_id, expert_id)) + + param = Parameter(torch.zeros(1)) + param.weight_loader = _wl + params = {"model.layers.1.mlp.experts.w13_weight": param} + dispatch = ExpertParamsDispatch( + mappings=(("experts.w13_", "experts.3.gate_proj.", 3, "w1"),) + ) + name = "model.layers.1.mlp.experts.3.gate_proj.weight" + target = dispatch.try_load(name, torch.ones(1), params) + self.assertEqual(target, "model.layers.1.mlp.experts.w13_weight") + self.assertEqual(calls[0][1], "w1") + + +class TestDeepseekV2WeightLoaderV2Gate(CustomTestCase): + def test_v2_env_dispatches_to_do_load_weights_v2(self): + from sglang.srt.models.deepseek_common.deepseek_weight_loader import ( + DeepseekV2WeightLoaderMixin, + ) + + class _Host(DeepseekV2WeightLoaderMixin): + pass + + host = _Host() + with patch.object(host, "_do_load_weights_v2", return_value={"x"}) as mock_v2: + with patch( + "sglang.srt.models.deepseek_common.deepseek_weight_loader.envs.SGLANG_ENABLE_WEIGHT_LOADER_V2.get", + return_value=True, + ): + result = host.do_load_weights([]) + mock_v2.assert_called_once() + self.assertEqual(result, {"x"}) + + +if __name__ == "__main__": + unittest.main()