diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 1694e651c773..e0d68bdfeb9d 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -207,6 +207,7 @@ class Envs: SGLANG_DISABLED_MODEL_ARCHS = EnvTuple(tuple()) SGLANG_PREFETCH_BLOCK_SIZE_MB = EnvInt(16) SGLANG_GEMMA_OUT_OF_PLACE_POSITION_MUTATION = EnvBool(False) + SGLANG_ENABLE_WEIGHT_LOADER_V2 = EnvBool(False) # 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..763b81288784 --- /dev/null +++ b/python/sglang/srt/model_loader/auto_loader.py @@ -0,0 +1,643 @@ +# 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.""" + +from __future__ import annotations + +import re + +from collections.abc import Callable, Iterable +from typing import Union + +import msgspec +import torch +from torch import nn +from torch.nn import Parameter + +from sglang.srt.layers.utils.common import get_layer_id +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.models.utils import AutoWeightsLoader, WeightsMapper + +__all__ = [ + "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", + + "MultiInputFusion", + "DEEPSEEK_GATE_UP_MAPPING", + "remap_fused_shared_expert_names", + "maybe_remap_deepseek_mla_kv_scale",] + + +class StackedParamsDispatch(msgspec.Struct, frozen=True): + mappings: tuple[tuple[str, str, Union[int, str]], ...] = () + + def try_load( + self, + name: str, + tensor: torch.Tensor, + params_dict: dict[str, Parameter], + ) -> str | None: + 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: + return target + param.weight_loader(param, tensor, shard_id) + return target + return None + + +STANDARD_QKV_MAPPING = StackedParamsDispatch( + mappings=( + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ) +) + +STANDARD_GATE_UP_MAPPING = StackedParamsDispatch( + mappings=( + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ) +) + +STANDARD_STACKED_MAPPING = StackedParamsDispatch( + mappings=( + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ) +) + +LLAMA_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), + ) +) + +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, + 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: + 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) + wl(params_dict[name], tensor) + loaded.add(name) + elif not any(name.endswith(suffix) for suffix in ignore_unexpected_suffixes): + pass + return loaded + + +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( + weights: Iterable[tuple[str, torch.Tensor]], + start_layer: int, + end_layer: int, +) -> Iterable[tuple[str, torch.Tensor]]: + 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 + + +_REMAP_REGISTRY: dict[str, Callable[[nn.Module], WeightsMapper]] = {} + + +def register_weight_remap(*class_names: str): + def decorator(fn: Callable[[nn.Module], WeightsMapper]): + for cn in class_names: + _REMAP_REGISTRY[cn] = fn + return fn + + return decorator + + +def get_weight_remap(model: nn.Module) -> WeightsMapper | None: + fn = _REMAP_REGISTRY.get(type(model).__name__) + if fn is None: + return None + return fn(model) + + +@register_weight_remap("LlamaForCausalLM") +def _llama_remap(model: nn.Module) -> WeightsMapper: + return WeightsMapper( + orig_to_new_suffix={ + ".activation_scale": ".input_scale", + ".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/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/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/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/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/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/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/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.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/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.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/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/manual/test_weight_loader_v2_equiv.py b/test/manual/test_weight_loader_v2_equiv.py new file mode 100644 index 000000000000..051df1e47616 --- /dev/null +++ b/test/manual/test_weight_loader_v2_equiv.py @@ -0,0 +1,138 @@ +# Copyright 2023-2025 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +# Manual verification for weight loader v2 (Qwen2 native path). +# +# Run: +# CUDA_VISIBLE_DEVICES=0 python test/manual/test_weight_loader_v2_equiv.py +# +# Engine-level e2e (Qwen2 + transformers backend) lives in: +# test/registered/model_loading/test_weight_loader_v2_e2e.py + +import unittest + +import torch + +from sglang.srt.environ import envs + +MODEL = "Qwen/Qwen2-0.5B" +QWEN3_MODEL = "Qwen/Qwen3-0.6B" + + +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: + 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 + 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_path, + dtype=torch.float16, + trust_remote_code=True, + ) + set_global_server_args_for_scheduler(server_args) + model_config = ModelConfig.from_server_args(server_args) + + with envs.SGLANG_ENABLE_WEIGHT_LOADER_V2.override(v2): + return get_model( + model_config=model_config, + load_config=LoadConfig(), + device_config=DeviceConfig(get_device()), + ) + + +def _state_dict_cpu(model: torch.nn.Module) -> dict[str, torch.Tensor]: + return { + name: param.detach().cpu().clone() for name, param in model.state_dict().items() + } + + +class TestWeightLoaderV2Equiv(unittest.TestCase): + @classmethod + def setUpClass(cls): + _init_model_parallel() + + @unittest.skipIf(not torch.cuda.is_available(), "needs GPU") + def test_qwen2_v1_v2_state_dict_identical(self): + model_v1 = _load_qwen2_native(v2=False) + state_v1 = _state_dict_cpu(model_v1) + del model_v1 + torch.cuda.empty_cache() + + model_v2 = _load_qwen2_native(v2=True) + state_v2 = _state_dict_cpu(model_v2) + del model_v2 + torch.cuda.empty_cache() + + self.assertEqual(set(state_v1.keys()), set(state_v2.keys())) + for name in sorted(state_v1.keys()): + torch.testing.assert_close( + state_v1[name], + state_v2[name], + rtol=0, + atol=0, + msg=name, + ) + + @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 new file mode 100644 index 000000000000..a43d6ef6c837 --- /dev/null +++ b/test/registered/model_loading/test_weight_loader_v2_e2e.py @@ -0,0 +1,109 @@ +# 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" +QWEN3_MODEL = "Qwen/Qwen3-0.6B" +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_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 + + 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/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() 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() 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()