model: serve bare Qwen3Model backbone natively as an embedding model (#32457)

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Jackey Hua
2026-07-27 15:49:58 +08:00
committed by GitHub
co-authored by Claude Opus 5
parent 169fc1e20c
commit 9a0bd24bed
4 changed files with 192 additions and 0 deletions
@@ -1696,6 +1696,7 @@ def is_generation_model(model_architectures: List[str], is_embedding: bool = Fal
or "Qwen3ForRewardModel" in model_architectures
or "Qwen2ForSequenceClassification" in model_architectures
or "Qwen3ForSequenceClassification" in model_architectures
or "Qwen3Model" in model_architectures
or "CLIPModel" in model_architectures
or "BertModel" in model_architectures
or "Contriever" in model_architectures
+124
View File
@@ -0,0 +1,124 @@
import logging
from typing import Iterable, Optional, Tuple
import torch
from torch import nn
from sglang.srt.layers.pooler import EmbeddingPoolerOutput, Pooler, PoolingType
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,
maybe_remap_kv_scale_name,
)
from sglang.srt.models.qwen3 import Qwen3Model as Qwen3TransformerModel
from sglang.srt.utils import add_prefix
logger = logging.getLogger(__name__)
class Qwen3Model(nn.Module):
"""Bare Qwen3 backbone (no LM head) served as an embedding model.
Checkpoints exported as architectures=["Qwen3Model"], e.g.
microsoft/harrier-oss-v1-0.6b, have no native implementation to resolve to
and fall back to the Transformers backend. Registering the arch here runs
them on the native fused Qwen3 kernels instead.
"""
def __init__(
self,
config,
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
) -> None:
super().__init__()
self.config = config
self.quant_config = quant_config
self.model = Qwen3TransformerModel(
config, quant_config=quant_config, prefix=add_prefix("model", prefix)
)
# Use LAST + normalize=True for qwen3 embedding based on official implementation
# Reference: https://github.com/QwenLM/Qwen3-Embedding/blob/main/examples/qwen3_embedding_transformers.py#L55
self.pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True)
def get_input_embeddings(self) -> nn.Embedding:
return self.model.get_input_embeddings()
@torch.no_grad()
def forward(
self,
input_ids: torch.Tensor,
positions: torch.Tensor,
forward_batch: ForwardBatch,
input_embeds: torch.Tensor = None,
get_embedding: bool = True,
) -> EmbeddingPoolerOutput:
assert get_embedding, f"{self.__class__.__name__} is only used for embedding"
hidden_states = self.model(input_ids, positions, forward_batch, input_embeds)
return self.pooler(hidden_states, forward_batch)
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
stacked_params_mapping = [
# (param_name, shard_name, shard_id)
("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),
]
params_dict = dict(self.named_parameters())
for name, loaded_weight in weights:
# Bare-backbone checkpoints omit the "model." prefix of the backbone
if not name.startswith("model.") and (
name.startswith("layers.")
or name.startswith("embed_tokens.")
or name.startswith("norm.")
):
name = add_prefix(name, "model")
# Skip rotary embeddings and other non-parameter tensors
if "rotary_emb.inv_freq" in name or "projector" in name:
continue
if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
# Models trained using ColossalAI may include these tensors in
# the checkpoint. Skip them.
continue
# Skip lm_head weights a non-tied checkpoint may carry (no LM head here)
if name.startswith("lm_head"):
continue
# Normalize kv cache scale names of quantized checkpoints
if "scale" in name:
name = maybe_remap_kv_scale_name(name, params_dict)
if name is None:
continue
for param_name, weight_name, shard_id in stacked_params_mapping:
if weight_name not in name:
continue
name = name.replace(weight_name, param_name)
if name.endswith(".bias") and name not in params_dict:
continue
param = params_dict[name]
weight_loader = param.weight_loader
weight_loader(param, loaded_weight, shard_id)
break
else:
if name.endswith(".bias") 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)
else:
logger.warning(f"Parameter {name} not found in params_dict")
EntryClass = Qwen3Model