diff --git a/docs_new/docs/supported-models/embedding_models.mdx b/docs_new/docs/supported-models/embedding_models.mdx index 3edf83449..86d3ad011 100644 --- a/docs_new/docs/supported-models/embedding_models.mdx +++ b/docs_new/docs/supported-models/embedding_models.mdx @@ -151,6 +151,12 @@ print("Embedding:", response["data"][0]["embedding"]) N/A Latest Qwen3-based text embedding model for semantic representation + + Qwen3 (bare backbone) + `microsoft/harrier-oss-v1-0.6b` + N/A + Bare Qwen3Model backbone (no LM head); served natively on SGLang's fused Qwen3 kernels and auto-classified as an embedding model + BGE `BAAI/bge-large-en-v1.5` diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 18417a394..992793c34 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -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 diff --git a/python/sglang/srt/models/qwen3_embedding.py b/python/sglang/srt/models/qwen3_embedding.py new file mode 100644 index 000000000..a32ee5d6d --- /dev/null +++ b/python/sglang/srt/models/qwen3_embedding.py @@ -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 diff --git a/test/registered/unit/models/test_qwen3_embedding_registration.py b/test/registered/unit/models/test_qwen3_embedding_registration.py new file mode 100644 index 000000000..8b2d82a52 --- /dev/null +++ b/test/registered/unit/models/test_qwen3_embedding_registration.py @@ -0,0 +1,61 @@ +"""Unit tests for native registration of the bare ``Qwen3Model`` embedding arch. + +Checkpoints such as ``microsoft/harrier-oss-v1-0.6b`` declare +``architectures=["Qwen3Model"]`` (a bare Qwen3 backbone). These must resolve to +the native SGLang implementation (``sglang.srt.models.qwen3_embedding.Qwen3Model``) +and be served as an embedding model, NOT fall back to the Transformers backend. +""" + +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=10, suite="base-a-test-cpu") + +import unittest + +from sglang.srt.configs.model_config import is_generation_model +from sglang.test.test_utils import CustomTestCase + + +class TestQwen3ModelEmbeddingRegistration(CustomTestCase): + def test_entry_class_is_native_qwen3model(self): + """The bare arch string maps to a native EntryClass named 'Qwen3Model'.""" + from sglang.srt.models import qwen3_embedding + + entry = qwen3_embedding.EntryClass + self.assertEqual(entry.__name__, "Qwen3Model") + self.assertEqual(entry.__module__, "sglang.srt.models.qwen3_embedding") + # It carries a LAST-token / normalized pooler, i.e. an embedding head. + self.assertTrue(hasattr(entry, "forward")) + self.assertTrue(hasattr(entry, "load_weights")) + + def test_registry_resolves_native_not_transformers_fallback(self): + """ModelRegistry resolves 'Qwen3Model' to the native class, not the + TransformersForCausalLM fallback.""" + from sglang.srt.models.registry import ModelRegistry + + model_cls, resolved_arch = ModelRegistry.resolve_model_cls("Qwen3Model") + self.assertEqual(resolved_arch, "Qwen3Model") + self.assertEqual(model_cls.__name__, "Qwen3Model") + self.assertEqual(model_cls.__module__, "sglang.srt.models.qwen3_embedding") + self.assertNotIn("Transformers", model_cls.__name__) + + def test_bare_qwen3model_classified_as_embedding(self): + """'Qwen3Model' is non-generative regardless of the --is-embedding flag.""" + self.assertFalse(is_generation_model(["Qwen3Model"])) + self.assertFalse(is_generation_model(["Qwen3Model"], is_embedding=False)) + self.assertFalse(is_generation_model(["Qwen3Model"], is_embedding=True)) + + def test_existing_qwen3_archs_unaffected(self): + """The generative / classification archs keep their prior behavior.""" + # Qwen3ForCausalLM is generative by default, embedding only with the flag + # (this is how Qwen3-Embedding-0.6B is served). + self.assertTrue(is_generation_model(["Qwen3ForCausalLM"])) + self.assertTrue(is_generation_model(["Qwen3ForCausalLM"], is_embedding=False)) + self.assertFalse(is_generation_model(["Qwen3ForCausalLM"], is_embedding=True)) + # Sequence-classification / reward archs stay non-generative. + self.assertFalse(is_generation_model(["Qwen3ForSequenceClassification"])) + self.assertFalse(is_generation_model(["Qwen3ForRewardModel"])) + + +if __name__ == "__main__": + unittest.main()