From 4f0409f8aa59baa26892712d5ad65fc6a4e1fc13 Mon Sep 17 00:00:00 2001 From: Shivam jindal <73436052+shvmjndl@users.noreply.github.com> Date: Mon, 16 Feb 2026 17:14:41 +0530 Subject: [PATCH] [Model] Add Qwen3ForRewardModel and fix Qwen3ForSequenceClassification (#17992) Co-authored-by: yes-its-shivam --- python/sglang/srt/configs/model_config.py | 1 + .../sglang/srt/models/qwen3_classification.py | 120 +++++++++++++----- python/sglang/srt/models/qwen3_rm.py | 47 +++++++ test/registered/models/test_reward_models.py | 2 + 4 files changed, 140 insertions(+), 30 deletions(-) create mode 100644 python/sglang/srt/models/qwen3_rm.py diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 701efb326..6fbd1db82 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -1192,6 +1192,7 @@ def is_generation_model(model_architectures: List[str], is_embedding: bool = Fal or "LlamaForSequenceClassificationWithNormal_Weights" in model_architectures or "InternLM2ForRewardModel" in model_architectures or "Qwen2ForRewardModel" in model_architectures + or "Qwen3ForRewardModel" in model_architectures or "Qwen2ForSequenceClassification" in model_architectures or "Qwen3ForSequenceClassification" in model_architectures or "CLIPModel" in model_architectures diff --git a/python/sglang/srt/models/qwen3_classification.py b/python/sglang/srt/models/qwen3_classification.py index a59d6769b..f6eff513d 100644 --- a/python/sglang/srt/models/qwen3_classification.py +++ b/python/sglang/srt/models/qwen3_classification.py @@ -12,6 +12,7 @@ # limitations under the License. # ============================================================================== +import logging from typing import Iterable, Optional, Tuple import torch @@ -21,11 +22,19 @@ from transformers import Qwen2Config # Qwen3 uses Qwen2Config 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.models.qwen3 import Qwen3ForCausalLM, Qwen3Model +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.models.qwen3 import Qwen3Model from sglang.srt.utils import add_prefix +logger = logging.getLogger(__name__) + + +class Qwen3ForPooledOutput(nn.Module): + """Base class for Qwen3 models that produce pooled output (classification, reward). + + Subclasses should set self.score and self.pooler in their __init__. + """ -class Qwen3ForSequenceClassification(nn.Module): def __init__( self, config: Qwen2Config, @@ -38,6 +47,85 @@ class Qwen3ForSequenceClassification(nn.Module): self.model = Qwen3Model( config, quant_config=quant_config, prefix=add_prefix("model", prefix) ) + self.eos_token_id = config.eos_token_id + # Subclasses must set self.score and self.pooler + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: Optional[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) + logits = self.score(hidden_states) + pooled_logits = self.pooler(logits, forward_batch).embeddings + + return EmbeddingPoolerOutput(pooled_logits) + + 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: + # Skip lm_head weights (pooled output models don't have lm_head) + if name.startswith("lm_head"): + continue + + # 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: + continue + + # Handle stacked parameters (qkv_proj, gate_up_proj) + 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) + # Skip loading extra bias for GPTQ models + if name.endswith(".bias") and name not in params_dict: + continue + if name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + # Skip loading extra bias for GPTQ models + 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") + + +class Qwen3ForSequenceClassification(Qwen3ForPooledOutput): + def __init__( + self, + config: Qwen2Config, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__(config, quant_config, prefix) self.score = nn.Linear(config.hidden_size, config.num_labels) # Use normalize=True for qwen3 embedding based on official implementation # Reference: https://github.com/QwenLM/Qwen3-Embedding/blob/main/examples/qwen3_embedding_transformers.py#L55 @@ -50,34 +138,6 @@ class Qwen3ForSequenceClassification(nn.Module): self.pooler = Pooler(pooling_type=PoolingType.LAST, normalize=normalize) - self.eos_token_id = config.eos_token_id - - @torch.no_grad() - def forward( - self, - input_ids: torch.Tensor, - positions: torch.Tensor, - forward_batch: ForwardBatch, - input_embeds: Optional[torch.Tensor] = None, - get_embedding: bool = True, - ) -> EmbeddingPoolerOutput: - assert ( - get_embedding - ), "Qwen3ForSequenceClassification is only used for embedding" - - hidden_states = self.model(input_ids, positions, forward_batch, input_embeds) - logits = self.score(hidden_states) - pooled_logits = self.pooler(logits, forward_batch).embeddings - - return EmbeddingPoolerOutput(pooled_logits) - - def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): - # Filter out lm_head weights of Qwen3ForCausalLM - filtered_weights = [ - (name, w) for name, w in weights if not name.startswith("lm_head") - ] - return Qwen3ForCausalLM.load_weights(self, filtered_weights) - EntryClass = [ Qwen3ForSequenceClassification, diff --git a/python/sglang/srt/models/qwen3_rm.py b/python/sglang/srt/models/qwen3_rm.py new file mode 100644 index 000000000..eeed15421 --- /dev/null +++ b/python/sglang/srt/models/qwen3_rm.py @@ -0,0 +1,47 @@ +# Copyright 2023-2024 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. +# ============================================================================== +"""Qwen3 Reward Model for RLHF and best-of-N sampling.""" + +from typing import Optional + +from torch import nn +from transformers import Qwen2Config # Qwen3 uses Qwen2Config + +from sglang.srt.layers.pooler import Pooler, PoolingType +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.models.qwen3_classification import Qwen3ForPooledOutput + + +class Qwen3ForRewardModel(Qwen3ForPooledOutput): + """Qwen3 Reward Model with 2-layer MLP scoring head for RLHF.""" + + def __init__( + self, + config: Qwen2Config, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__(config, quant_config, prefix) + self.num_labels = 1 + self.score = nn.Sequential( + nn.Linear(config.hidden_size, config.hidden_size), + nn.ReLU(), + nn.Linear(config.hidden_size, self.num_labels), + ) + self.pooler = Pooler(pooling_type=PoolingType.LAST, normalize=False) + + +EntryClass = [ + Qwen3ForRewardModel, +] diff --git a/test/registered/models/test_reward_models.py b/test/registered/models/test_reward_models.py index 553530708..73fd5535f 100644 --- a/test/registered/models/test_reward_models.py +++ b/test/registered/models/test_reward_models.py @@ -30,6 +30,8 @@ register_amd_ci(est_time=132, suite="stage-b-test-small-1-gpu-amd") MODELS = [ ("LxzGordon/URM-LLaMa-3.1-8B", 1, 4e-2), ("Skywork/Skywork-Reward-Llama-3.1-8B-v0.2", 1, 4e-2), + # Qwen3-based reward model (uses Qwen3ForSequenceClassification) + ("Skywork/Skywork-Reward-V2-Qwen3-0.6B", 1, 1.5e-1), ] TORCH_DTYPES = [torch.float16]