226 lines
7.4 KiB
Python
226 lines
7.4 KiB
Python
# Copyright 2023-2025 SGLang Team
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ==============================================================================
|
|
"""Centralized weight loading utilities for native SGLang models.
|
|
|
|
This module provides:
|
|
- StackedParamsDispatch: reusable stacked-parameter routing (qkv_proj, gate_up_proj).
|
|
- ExpertParamsDispatch: MoE expert_id + w1/w2/w3 shard routing.
|
|
- filter_pp_weights: generator that drops out-of-range PP layers.
|
|
- RemapRegistry: architecture-specific name remap registration.
|
|
- Re-exports of AutoWeightsLoader and WeightsMapper from models/utils.py.
|
|
|
|
Load / post-load split (PR1 protocol, see #31051):
|
|
load_weights(..., run_post_load=True) -> WeightLoadResult
|
|
post_load_weights(loaded=result) -> optional GPU derivations (MLA w_kc/w_vc, etc.)
|
|
|
|
Migration plan: https://github.com/sgl-project/sglang/issues/31051 (RFC #24703).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Callable, Iterable
|
|
from typing import Union
|
|
|
|
import msgspec
|
|
import torch
|
|
from torch import nn
|
|
from torch.nn import Parameter
|
|
|
|
from sglang.srt.layers.utils.common import get_layer_id
|
|
from sglang.srt.models.utils import AutoWeightsLoader, WeightsMapper
|
|
|
|
__all__ = [
|
|
"AutoWeightsLoader",
|
|
"WeightsMapper",
|
|
"StackedParamsDispatch",
|
|
"STANDARD_QKV_MAPPING",
|
|
"STANDARD_GATE_UP_MAPPING",
|
|
"STANDARD_STACKED_MAPPING",
|
|
"LLAMA_STACKED_MAPPING",
|
|
"filter_pp_weights",
|
|
"register_weight_remap",
|
|
"get_weight_remap",
|
|
]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Stacked Parameters Dispatch
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class StackedParamsDispatch(msgspec.Struct, frozen=True):
|
|
"""Centralized stacked-parameter loading for fused linear layers.
|
|
|
|
Handles the common pattern of mapping checkpoint names
|
|
(q_proj, k_proj, v_proj, gate_proj, up_proj) to fused runtime parameters
|
|
(qkv_proj, gate_up_proj) with the correct shard IDs.
|
|
|
|
Quantization is handled entirely by ``param.weight_loader`` on the layer —
|
|
this class only routes the tensor to the correct parameter with the correct
|
|
shard_id.
|
|
|
|
Usage::
|
|
|
|
mapping = StackedParamsDispatch([
|
|
("qkv_proj", "q_proj", "q"),
|
|
("qkv_proj", "k_proj", "k"),
|
|
("qkv_proj", "v_proj", "v"),
|
|
])
|
|
target = mapping.try_load(name, tensor, params_dict)
|
|
"""
|
|
|
|
# (fused_param_name, checkpoint_source_name, shard_id).
|
|
mappings: tuple[tuple[str, str, Union[int, str]], ...] = ()
|
|
|
|
def try_load(
|
|
self,
|
|
name: str,
|
|
tensor: torch.Tensor,
|
|
params_dict: dict[str, Parameter],
|
|
) -> str | None:
|
|
"""Try to load a weight via stacked mapping.
|
|
|
|
Returns the loaded runtime parameter name if matched and loaded,
|
|
the target name (for skip tracking) if the target param is missing
|
|
(e.g. optional bias), or None if no mapping matched.
|
|
"""
|
|
for fused_name, source_name, shard_id in self.mappings:
|
|
if source_name not in name:
|
|
continue
|
|
target = name.replace(source_name, fused_name)
|
|
param = params_dict.get(target)
|
|
if param is None:
|
|
# Parameter doesn't exist — e.g. GPTQ bias.
|
|
# Return target so caller can track the skip.
|
|
return target
|
|
param.weight_loader(param, tensor, shard_id)
|
|
return target
|
|
return None
|
|
|
|
|
|
# Pre-built instances for the most common decoder patterns.
|
|
|
|
STANDARD_QKV_MAPPING = StackedParamsDispatch(
|
|
mappings=(
|
|
("qkv_proj", "q_proj", "q"),
|
|
("qkv_proj", "k_proj", "k"),
|
|
("qkv_proj", "v_proj", "v"),
|
|
)
|
|
)
|
|
|
|
STANDARD_GATE_UP_MAPPING = StackedParamsDispatch(
|
|
mappings=(
|
|
("gate_up_proj", "gate_proj", 0),
|
|
("gate_up_proj", "up_proj", 1),
|
|
)
|
|
)
|
|
|
|
STANDARD_STACKED_MAPPING = StackedParamsDispatch(
|
|
mappings=(
|
|
("qkv_proj", "q_proj", "q"),
|
|
("qkv_proj", "k_proj", "k"),
|
|
("qkv_proj", "v_proj", "v"),
|
|
("gate_up_proj", "gate_proj", 0),
|
|
("gate_up_proj", "up_proj", 1),
|
|
)
|
|
)
|
|
|
|
# Llama-family full-path stacked mapping (dot-prefixed shard names).
|
|
LLAMA_STACKED_MAPPING = StackedParamsDispatch(
|
|
mappings=(
|
|
(".qkv_proj", ".q_proj", "q"),
|
|
(".qkv_proj", ".k_proj", "k"),
|
|
(".qkv_proj", ".v_proj", "v"),
|
|
(".gate_up_proj", ".gate_proj", 0),
|
|
(".gate_up_proj", ".up_proj", 1),
|
|
)
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Pipeline Parallel Weight Filter
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def filter_pp_weights(
|
|
weights: Iterable[tuple[str, torch.Tensor]],
|
|
start_layer: int,
|
|
end_layer: int,
|
|
) -> Iterable[tuple[str, torch.Tensor]]:
|
|
"""Drop checkpoint entries whose layer index is outside [start_layer, end_layer).
|
|
|
|
Weights that don't contain a parseable layer index (embed_tokens, lm_head,
|
|
layer norms, etc.) are always passed through.
|
|
"""
|
|
for name, tensor in weights:
|
|
layer_id = get_layer_id(name)
|
|
if layer_id is not None and (layer_id < start_layer or layer_id >= end_layer):
|
|
continue
|
|
yield name, tensor
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Weight Remap Registry
|
|
# ---------------------------------------------------------------------------
|
|
|
|
_REMAP_REGISTRY: dict[str, Callable[[nn.Module], WeightsMapper]] = {}
|
|
|
|
|
|
def register_weight_remap(*class_names: str):
|
|
"""Decorator to register an architecture-specific weight remap function.
|
|
|
|
The decorated function receives a model instance and returns a
|
|
``WeightsMapper``. If no remap is needed for a model, do not register it.
|
|
|
|
Example::
|
|
|
|
@register_weight_remap("LlamaForCausalLM")
|
|
def _llama_remap(model) -> WeightsMapper:
|
|
return WeightsMapper(orig_to_new_suffix={
|
|
".activation_scale": ".input_scale",
|
|
".weight_scale_inv": ".weight_scale",
|
|
})
|
|
"""
|
|
|
|
def decorator(fn: Callable[[nn.Module], WeightsMapper]):
|
|
for cn in class_names:
|
|
_REMAP_REGISTRY[cn] = fn
|
|
return fn
|
|
|
|
return decorator
|
|
|
|
|
|
def get_weight_remap(model: nn.Module) -> WeightsMapper | None:
|
|
"""Get the registered weight remap for a model instance, or None."""
|
|
fn = _REMAP_REGISTRY.get(type(model).__name__)
|
|
if fn is None:
|
|
return None
|
|
return fn(model)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Architecture-Specific Registrations
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@register_weight_remap("LlamaForCausalLM")
|
|
def _llama_remap(model: nn.Module) -> WeightsMapper:
|
|
"""Llama-family FP8 scale suffix normalization."""
|
|
return WeightsMapper(
|
|
orig_to_new_suffix={
|
|
".activation_scale": ".input_scale",
|
|
".weight_scale_inv": ".weight_scale",
|
|
}
|
|
)
|