@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 114 lines 4 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 # Adapted from
5 # https://huggingface.co/Qwen/Qwen2.5-Math-RM-72B/blob/main/modeling_qwen2_rm.py
6 # Copyright 2024 The Qwen team.
7 # Copyright 2023 The vLLM team.
8 """Inference-only Qwen2-RM model compatible with HuggingFace weights."""
9
10 from collections.abc import Iterable
11
12 import torch
13 from torch import nn
14
15 from vllm.config import VllmConfig
16 from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear
17 from vllm.model_executor.layers.pooler import Pooler
18 from vllm.model_executor.layers.pooler.tokwise import pooler_for_token_classify
19 from vllm.sequence import IntermediateTensors
20
21 from .interfaces import SupportsLoRA, SupportsPP
22 from .interfaces_base import default_pooling_type
23 from .qwen2 import Qwen2Model
24 from .utils import AutoWeightsLoader, maybe_prefix
25
26
27 class Qwen2RewardBaseModel(nn.Module, SupportsLoRA, SupportsPP):
28 is_pooling_model = True
29 pooler: Pooler
30
31 hf_to_vllm_mapper = Qwen2Model.hf_to_vllm_mapper
32 packed_modules_mapping = {
33 "qkv_proj": ["q_proj", "k_proj", "v_proj"],
34 "gate_up_proj": ["gate_proj", "up_proj"],
35 }
36
37 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
38 super().__init__()
39 config = vllm_config.model_config.hf_config
40 quant_config = vllm_config.quant_config
41
42 self.config = config
43
44 self.quant_config = quant_config
45 self.model = Qwen2Model(
46 vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
47 )
48 self.head_dtype = vllm_config.model_config.head_dtype
49
50 self.score = nn.Sequential(
51 ColumnParallelLinear(
52 config.hidden_size,
53 config.hidden_size,
54 quant_config=quant_config,
55 params_dtype=self.head_dtype,
56 return_bias=False,
57 ),
58 nn.ReLU(),
59 RowParallelLinear(
60 config.hidden_size,
61 config.num_labels,
62 params_dtype=self.head_dtype,
63 quant_config=quant_config,
64 return_bias=False,
65 ),
66 )
67 self.make_empty_intermediate_tensors = (
68 self.model.make_empty_intermediate_tensors
69 )
70
71 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
72 return self.model.embed_input_ids(input_ids)
73
74 def forward(
75 self,
76 input_ids: torch.Tensor | None,
77 positions: torch.Tensor,
78 intermediate_tensors: IntermediateTensors | None = None,
79 inputs_embeds: torch.Tensor | None = None,
80 ) -> torch.Tensor | IntermediateTensors:
81 hidden_states = self.model(
82 input_ids, positions, intermediate_tensors, inputs_embeds
83 )
84 hidden_states = hidden_states.to(self.head_dtype)
85 logits = self.score(hidden_states)
86 return logits
87
88 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
89 loader = AutoWeightsLoader(self, ignore_unexpected_prefixes=["lm_head."])
90 return loader.load_weights(weights)
91
92
93 @default_pooling_type(tok_pooling_type="ALL")
94 class Qwen2ForRewardModel(Qwen2RewardBaseModel):
95 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
96 vllm_config.model_config.hf_config.num_labels = 1
97 super().__init__(vllm_config=vllm_config, prefix=prefix)
98
99 pooler_config = vllm_config.model_config.pooler_config
100 assert pooler_config is not None
101
102 self.pooler = pooler_for_token_classify(pooler_config)
103
104
105 @default_pooling_type(tok_pooling_type="STEP")
106 class Qwen2ForProcessRewardModel(Qwen2RewardBaseModel):
107 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
108 vllm_config.model_config.hf_config.num_labels = 2
109 super().__init__(vllm_config=vllm_config, prefix=prefix)
110
111 pooler_config = vllm_config.model_config.pooler_config
112 assert pooler_config is not None
113
114 self.pooler = pooler_for_token_classify(pooler_config)