@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 6e4efafc6e

[Model] Migrate GPT-NeoX, Phi, Seed-OSS and Jais2 to the Transformers modeling backend (#59701)

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

Harry Mellor committed Oct 2, 2026 at 12:18 UTC 6e4efafc6edfee8be31c3c147164a1daa4650051
8 files changed +26 -1536
docs/models/supported_models.md
+6 -6
@@ -410,7 +410,6 @@ th {
410 | `GlmMoeDsaForCausalLM` | GLM-5, GLM-5.1, GLM-5.2 | `zai-org/GLM-5`, etc. | ✅︎ | ✅︎ |
411 | `GPT2LMHeadModel` | GPT-2 | `openai-community/gpt2`, `openai-community/gpt2-xl`, etc. | | ✅︎ |
412 | `GPTJForCausalLM` | GPT-J | `EleutherAI/gpt-j-6b`, `nomic-ai/gpt4all-j`, etc. | | ✅︎ |
413 -| `GPTNeoXForCausalLM` | GPT-NeoX, Pythia, OpenAssistant, Dolly V2, StableLM | `EleutherAI/gpt-neox-20b`, `EleutherAI/pythia-12b`, `OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5`, `databricks/dolly-v2-12b`, `stabilityai/stablelm-tuned-alpha-7b`, etc. | | ✅︎ |
413 | `GptOssForCausalLM` | GPT-OSS | `openai/gpt-oss-120b`, `openai/gpt-oss-20b` | ✅︎ | ✅︎ |
414 | `GraniteForCausalLM` | Granite 3.0, Granite 3.1, PowerLM | `ibm-granite/granite-3.0-2b-base`, `ibm-granite/granite-3.1-8b-instruct`, `ibm/PowerLM-3b`, etc. | ✅︎ | ✅︎ |
415 | `GraniteMoeForCausalLM` | Granite 3.0 MoE, PowerMoE | `ibm-granite/granite-3.0-1b-a400m-base`, `ibm-granite/granite-3.0-3b-a800m-instruct`, `ibm/PowerMoE-3b`, etc. | ✅︎ | ✅︎ |
@@ -426,7 +425,6 @@ th {
425 | `InternLM3ForCausalLM` | InternLM3 | `internlm/internlm3-8b-instruct`, etc. | ✅︎ | ✅︎ |
426 | `IQuestCoderForCausalLM` | IQuestCoderV1 | `IQuestLab/IQuest-Coder-V1-40B-Instruct`, etc. | | |
427 | `IQuestLoopCoderForCausalLM` | IQuestLoopCoderV1 | `IQuestLab/IQuest-Coder-V1-40B-Loop-Instruct`, etc. | | |
429 -| `Jais2ForCausalLM` | Jais2 | `inceptionai/Jais-2-8B-Chat`, `inceptionai/Jais-2-70B-Chat`, etc. | | ✅︎ |
428 | `JambaForCausalLM` | Jamba | `ai21labs/AI21-Jamba-1.5-Large`, `ai21labs/AI21-Jamba-1.5-Mini`, `ai21labs/Jamba-v0.1`, etc. | ✅︎ | ✅︎ |
429 | `K2HorizonForCausalLM` | K2-Horizon | `IFM/K2-Horizon-7B`, `IFM/K2-Horizon-375B` | | ✅︎ |
430 | `KimiLinearForCausalLM` | Kimi-Linear-48B-A3B-Base, Kimi-Linear-48B-A3B-Instruct | `moonshotai/Kimi-Linear-48B-A3B-Base`, `moonshotai/Kimi-Linear-48B-A3B-Instruct` | | ✅︎ |
@@ -456,7 +454,6 @@ th {
454 | `PanguProMoEV2ForCausalLM` | openpangu-pro-moe-v2 | | ✅︎ | ✅︎ |
455 | `PanguUltraMoEForCausalLM` | openpangu-ultra-moe-718b-model | `FreedomIntelligence/openPangu-Ultra-MoE-718B-V1.1` | ✅︎ | ✅︎ |
456 | `Param2MoEForCausalLM` | param2moe | `bharatgenai/Param2-17B-A2.4B-Thinking`, etc. | ✅︎ | ✅︎ |
459 -| `PhiForCausalLM` | Phi | `microsoft/phi-1_5`, `microsoft/phi-2`, etc. | ✅︎ | ✅︎ |
457 | `Phi3ForCausalLM` | Phi-4, Phi-3 | `microsoft/Phi-4-mini-instruct`, `microsoft/Phi-4`, `microsoft/Phi-3-mini-4k-instruct`, `microsoft/Phi-3-mini-128k-instruct`, `microsoft/Phi-3-medium-128k-instruct`, etc. | ✅︎ | ✅︎ |
458 | `PhiMoEForCausalLM` | Phi-3.5-MoE | `microsoft/Phi-3.5-MoE-instruct`, etc. | ✅︎ | ✅︎ |
459 | `Plamo3ForCausalLM` | PLaMo3 | `pfnet/plamo-3-nict-2b-base`, `pfnet/plamo-3-nict-8b-base`, etc. | ✅ | ✅︎ |
@@ -468,7 +465,6 @@ th {
465 | `Rnj1ForCausalLM` | Rnj1 | `EssentialAI/rnj-1-instruct`, etc. | | |
466 | `SarvamMoEForCausalLM` | Sarvam 2 | `sarvamai/sarvam2-30b-a3b`, etc. | ✅︎ | ✅︎ |
467 | `SarvamMLAForCausalLM` | Sarvam 2 | `sarvamai/sarvam2-105b-a9b`, etc. | | ✅︎ |
471 -| `SeedOssForCausalLM` | SeedOss | `ByteDance-Seed/Seed-OSS-36B-Instruct`, etc. | ✅︎ | ✅︎ |
468 | `SolarForCausalLM` | Solar Pro | `upstage/solar-pro-preview-instruct`, etc. | ✅︎ | ✅︎ |
469 | `StableLmForCausalLM` | StableLM | `stabilityai/stablelm-3b-4e1t`, `stabilityai/stablelm-base-alpha-7b-v2`, etc. | | |
470 | `Step1ForCausalLM` | Step-Audio | `stepfun-ai/Step-Audio-EditX`, etc. | ✅︎ | ✅︎ |
@@ -484,16 +480,20 @@ Some models are supported only via the [Transformers modeling backend](#transfor
480 | ------------ | ------ | ----------------- | -------------------- | ------------------------- |
481 | `ArceeForCausalLM` | Arcee (AFM) | `arcee-ai/AFM-4.5B-Base`, etc. | ✅︎ | ✅︎ |
482 | `CwmForCausalLM` | CWM | `facebook/cwm`, etc. | ✅︎ | ✅︎ |
487 -| `FlexOlmoForCausalLM` | FlexOlmo | `allenai/FlexOlmo-7x7B-1T`, `allenai/FlexOlmo-7x7B-1T-RT`, etc. | | ✅︎ |
483 +| `FlexOlmoForCausalLM` | FlexOlmo | `allenai/FlexOlmo-7x7B-1T`, `allenai/FlexOlmo-7x7B-1T-RT`, etc. | ✅︎ | ✅︎ |
484 | `GlmForCausalLM` | GLM-4 | `zai-org/glm-4-9b-chat-hf`, etc. | ✅︎ | ✅︎ |
485 | `GPTBigCodeForCausalLM` | StarCoder, SantaCoder, WizardCoder | `bigcode/starcoder`, `bigcode/gpt_bigcode-santacoder`, `WizardLM/WizardCoder-15B-V1.0`, etc. | ✅︎ | |
486 +| `GPTNeoXForCausalLM` | GPT-NeoX, Pythia, OpenAssistant, Dolly V2, StableLM | `EleutherAI/gpt-neox-20b`, `EleutherAI/pythia-12b`, `OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5`, `databricks/dolly-v2-12b`, `stabilityai/stablelm-tuned-alpha-7b`, etc. | | ✅︎ |
487 | `HunYuanDenseV1ForCausalLM` | Hunyuan Dense | `tencent/Hunyuan-7B-Instruct` | ✅︎ | ✅︎ |
488 | `HunYuanMoEV1ForCausalLM` | Hunyuan-A13B | `tencent/Hunyuan-A13B-Instruct`, `tencent/Hunyuan-A13B-Pretrain`, `tencent/Hunyuan-A13B-Instruct-FP8`, etc. | ✅︎ | ✅︎ |
492 -| `MellumForCausalLM` | Mellum 2 | `JetBrains/Mellum2-12B-A2.5B-Base`, etc. | | ✅︎ |
489 +| `Jais2ForCausalLM` | Jais2 | `inceptionai/Jais-2-8B-Chat`, `inceptionai/Jais-2-70B-Chat`, etc. | ✅︎ | ✅︎ |
490 +| `MellumForCausalLM` | Mellum 2 | `JetBrains/Mellum2-12B-A2.5B-Base`, etc. | ✅︎ | ✅︎ |
491 | `NanbeigeForCausalLM` | Nanbeige4.2 | `Nanbeige/Nanbeige4.2-3B`, etc. | ✅︎ | ✅︎ |
492 | `OlmoForCausalLM` | OLMo | `allenai/OLMo-1B-hf`, `allenai/OLMo-7B-hf`, etc. | ✅︎ | ✅︎ |
493 | `Olmo2ForCausalLM` | OLMo2 | `allenai/OLMo-2-0425-1B`, etc. | ✅︎ | ✅︎ |
494 | `Olmo3ForCausalLM` | OLMo3 | `allenai/Olmo-3-7B-Instruct`, `allenai/Olmo-3-32B-Think`, etc. | ✅︎ | ✅︎ |
495 +| `PhiForCausalLM` | Phi | `microsoft/phi-1_5`, `microsoft/phi-2`, etc. | ✅︎ | ✅︎ |
496 +| `SeedOssForCausalLM` | SeedOss | `ByteDance-Seed/Seed-OSS-36B-Instruct`, etc. | ✅︎ | ✅︎ |
497 | `SmolLM3ForCausalLM` | SmolLM3 | `HuggingFaceTB/SmolLM3-3B` | ✅︎ | ✅︎ |
498 | `Starcoder2ForCausalLM` | Starcoder2 | `bigcode/starcoder2-3b`, `bigcode/starcoder2-7b`, `bigcode/starcoder2-15b`, etc. | ✅︎ | ✅︎ |
499 | `VaultGemmaForCausalLM` | VaultGemma | `google/vaultgemma-1b` | ✅︎ | ✅︎ |
vllm/model_executor/models/gpt_neox.py deleted
-324
@@ -1,324 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -# Adapted from
5 -# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/gpt_neox/modeling_gpt_neox.py
6 -# Copyright 2023 The vLLM team.
7 -# Copyright 2022 EleutherAI The HuggingFace Inc. team. All rights reserved.
8 -#
9 -# Licensed under the Apache License, Version 2.0 (the "License");
10 -# you may not use this file except in compliance with the License.
11 -# You may obtain a copy of the License at
12 -#
13 -# http://www.apache.org/licenses/LICENSE-2.0
14 -#
15 -# Unless required by applicable law or agreed to in writing, software
16 -# distributed under the License is distributed on an "AS IS" BASIS,
17 -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
18 -# See the License for the specific language governing permissions and
19 -# limitations under the License.
20 -"""Inference-only GPT-NeoX model compatible with HuggingFace weights."""
21 -
22 -from collections.abc import Iterable
23 -from itertools import islice
24 -
25 -import torch
26 -from torch import nn
27 -from transformers import GPTNeoXConfig
28 -
29 -from vllm.compilation.decorators import support_torch_compile
30 -from vllm.config import CacheConfig, VllmConfig
31 -from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
32 -from vllm.model_executor.layers.activation import get_act_fn
33 -from vllm.model_executor.layers.attention import Attention
34 -from vllm.model_executor.layers.linear import (
35 - ColumnParallelLinear,
36 - QKVParallelLinear,
37 - RowParallelLinear,
38 -)
39 -from vllm.model_executor.layers.logits_processor import LogitsProcessor
40 -from vllm.model_executor.layers.quantization import QuantizationConfig
41 -from vllm.model_executor.layers.rotary_embedding import get_rope
42 -from vllm.model_executor.layers.vocab_parallel_embedding import (
43 - ParallelLMHead,
44 - VocabParallelEmbedding,
45 -)
46 -from vllm.sequence import IntermediateTensors
47 -
48 -from .interfaces import SupportsPP
49 -from .utils import (
50 - AutoWeightsLoader,
51 - WeightsMapper,
52 - make_empty_intermediate_tensors_factory,
53 - make_layers,
54 - maybe_prefix,
55 -)
56 -
57 -
58 -class GPTNeoXAttention(nn.Module):
59 - def __init__(
60 - self,
61 - config: GPTNeoXConfig,
62 - cache_config: CacheConfig | None = None,
63 - quant_config: QuantizationConfig | None = None,
64 - prefix: str = "",
65 - ):
66 - super().__init__()
67 - self.total_num_heads = config.num_attention_heads
68 - self.hidden_size = config.hidden_size
69 - self.head_size = self.hidden_size // self.total_num_heads
70 - self.bias = getattr(config, "attention_bias", True)
71 -
72 - tensor_model_parallel_world_size = get_tensor_model_parallel_world_size()
73 - assert self.total_num_heads % tensor_model_parallel_world_size == 0
74 - self.num_heads = self.total_num_heads // tensor_model_parallel_world_size
75 -
76 - self.query_key_value = QKVParallelLinear(
77 - config.hidden_size,
78 - self.head_size,
79 - self.total_num_heads,
80 - bias=self.bias,
81 - quant_config=quant_config,
82 - prefix=f"{prefix}.query_key_value",
83 - )
84 - self.dense = RowParallelLinear(
85 - config.hidden_size,
86 - config.hidden_size,
87 - bias=self.bias,
88 - quant_config=quant_config,
89 - prefix=f"{prefix}.dense",
90 - )
91 - max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
92 - self.rotary_emb = get_rope(
93 - self.head_size,
94 - max_position=max_position_embeddings,
95 - rope_parameters=config.rope_parameters,
96 - )
97 - scaling = self.head_size**-0.5
98 - self.attn = Attention(
99 - self.num_heads,
100 - self.head_size,
101 - scaling,
102 - cache_config=cache_config,
103 - quant_config=quant_config,
104 - prefix=f"{prefix}.attn",
105 - )
106 -
107 - def forward(
108 - self,
109 - position_ids: torch.Tensor,
110 - hidden_states: torch.Tensor,
111 - ) -> torch.Tensor:
112 - qkv, _ = self.query_key_value(hidden_states)
113 - q, k, v = qkv.chunk(chunks=3, dim=-1)
114 - q, k = self.rotary_emb(position_ids, q, k)
115 - attn_output = self.attn(q, k, v)
116 - output, _ = self.dense(attn_output)
117 - return output
118 -
119 -
120 -class GPTNeoXMLP(nn.Module):
121 - def __init__(
122 - self,
123 - config: GPTNeoXConfig,
124 - quant_config: QuantizationConfig | None = None,
125 - prefix: str = "",
126 - ):
127 - super().__init__()
128 - self.dense_h_to_4h = ColumnParallelLinear(
129 - config.hidden_size,
130 - config.intermediate_size,
131 - quant_config=quant_config,
132 - prefix=f"{prefix}.dense_h_to_4h",
133 - )
134 - self.dense_4h_to_h = RowParallelLinear(
135 - config.intermediate_size,
136 - config.hidden_size,
137 - quant_config=quant_config,
138 - prefix=f"{prefix}.dense_4h_to_h",
139 - )
140 - self.act = get_act_fn(config.hidden_act)
141 -
142 - def forward(self, hidden_states):
143 - hidden_states, _ = self.dense_h_to_4h(hidden_states)
144 - hidden_states = self.act(hidden_states)
145 - hidden_states, _ = self.dense_4h_to_h(hidden_states)
146 - return hidden_states
147 -
148 -
149 -class GPTNeoXLayer(nn.Module):
150 - def __init__(
151 - self,
152 - config: GPTNeoXConfig,
153 - cache_config: CacheConfig | None = None,
154 - quant_config: QuantizationConfig | None = None,
155 - prefix: str = "",
156 - ):
157 - super().__init__()
158 - self.use_parallel_residual = config.use_parallel_residual
159 - self.input_layernorm = nn.LayerNorm(
160 - config.hidden_size, eps=config.layer_norm_eps
161 - )
162 - self.post_attention_layernorm = nn.LayerNorm(
163 - config.hidden_size, eps=config.layer_norm_eps
164 - )
165 - self.attention = GPTNeoXAttention(
166 - config, cache_config, quant_config, prefix=f"{prefix}.attention"
167 - )
168 - self.mlp = GPTNeoXMLP(config, quant_config, prefix=f"{prefix}.mlp")
169 -
170 - def forward(
171 - self,
172 - position_ids: torch.Tensor,
173 - hidden_states: torch.Tensor,
174 - ) -> torch.Tensor:
175 - attn_input = self.input_layernorm(hidden_states)
176 - attn_output = self.attention(
177 - position_ids=position_ids,
178 - hidden_states=attn_input,
179 - )
180 -
181 - if self.use_parallel_residual:
182 - # pseudocode:
183 - # x = x + attn(ln1(x)) + mlp(ln2(x))
184 - mlp_input = self.post_attention_layernorm(hidden_states)
185 - mlp_output = self.mlp(mlp_input)
186 - hidden_states = mlp_output + attn_output + hidden_states
187 - else:
188 - # pseudocode:
189 - # x = x + attn(ln1(x))
190 - # x = x + mlp(ln2(x))
191 - attn_output = attn_output + hidden_states
192 - mlp_input = self.post_attention_layernorm(attn_output)
193 - mlp_output = self.mlp(mlp_input)
194 - hidden_states = mlp_output + attn_output
195 - return hidden_states
196 -
197 -
198 -@support_torch_compile
199 -class GPTNeoXModel(nn.Module):
200 - hf_to_vllm_mapper = WeightsMapper(
201 - orig_to_new_substr={"attention.bias": None, "attention.masked_bias": None}
202 - )
203 -
204 - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
205 - super().__init__()
206 -
207 - config = vllm_config.model_config.hf_config
208 - cache_config = vllm_config.cache_config
209 - quant_config = vllm_config.quant_config
210 -
211 - self.config = config
212 -
213 - self.embed_in = VocabParallelEmbedding(
214 - config.vocab_size,
215 - config.hidden_size,
216 - )
217 - self.start_layer, self.end_layer, self.layers = make_layers(
218 - config.num_hidden_layers,
219 - lambda prefix: GPTNeoXLayer(
220 - config, cache_config, quant_config, prefix=prefix
221 - ),
222 - prefix=f"{prefix}.layers",
223 - )
224 - self.final_layer_norm = nn.LayerNorm(
225 - config.hidden_size, eps=config.layer_norm_eps
226 - )
227 - self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
228 - ["hidden_states"], config.hidden_size
229 - )
230 -
231 - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
232 - return self.embed_in(input_ids)
233 -
234 - def forward(
235 - self,
236 - input_ids: torch.Tensor | None,
237 - position_ids: torch.Tensor,
238 - intermediate_tensors: IntermediateTensors | None,
239 - inputs_embeds: torch.Tensor | None = None,
240 - ) -> torch.Tensor | IntermediateTensors:
241 - if get_pp_group().is_first_rank:
242 - if inputs_embeds is not None:
243 - hidden_states = inputs_embeds
244 - else:
245 - hidden_states = self.embed_input_ids(input_ids)
246 - else:
247 - assert intermediate_tensors is not None
248 - hidden_states = intermediate_tensors["hidden_states"]
249 - for layer in islice(self.layers, self.start_layer, self.end_layer):
250 - hidden_states = layer(position_ids, hidden_states)
251 - if not get_pp_group().is_last_rank:
252 - return IntermediateTensors({"hidden_states": hidden_states})
253 - hidden_states = self.final_layer_norm(hidden_states)
254 - return hidden_states
255 -
256 - def _repack_qkv(
257 - self, weights: Iterable[tuple[str, torch.Tensor]]
258 - ) -> Iterable[tuple[str, torch.Tensor]]:
259 - # GPT-NeoX's fused QKV is laid out as (num_heads * 3 * head_size) on
260 - # its output dim (0), while vLLM expects (3 * num_heads * head_size).
261 - num_heads = self.config.num_attention_heads
262 - for name, loaded_weight in weights:
263 - if "query_key_value" in name:
264 - shape = loaded_weight.shape
265 - loaded_weight = loaded_weight.view((num_heads, 3, -1) + shape[1:])
266 - loaded_weight = loaded_weight.transpose(0, 1)
267 - loaded_weight = loaded_weight.reshape(shape)
268 - yield name, loaded_weight
269 -
270 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
271 - loader = AutoWeightsLoader(self)
272 - return loader.load_weights(
273 - self._repack_qkv(weights), mapper=self.hf_to_vllm_mapper
274 - )
275 -
276 -
277 -class GPTNeoXForCausalLM(nn.Module, SupportsPP):
278 - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
279 - super().__init__()
280 - config = vllm_config.model_config.hf_config
281 - quant_config = vllm_config.quant_config
282 - self.config = config
283 - self.quant_config = quant_config
284 - self.gpt_neox = GPTNeoXModel(
285 - vllm_config=vllm_config, prefix=maybe_prefix(prefix, "gpt_neox")
286 - )
287 - self.embed_out = ParallelLMHead(
288 - config.vocab_size,
289 - config.hidden_size,
290 - quant_config=quant_config,
291 - prefix=maybe_prefix(prefix, "embed_out"),
292 - )
293 - if self.config.tie_word_embeddings:
294 - self.embed_out = self.embed_out.tie_weights(self.gpt_neox.embed_in)
295 - self.logits_processor = LogitsProcessor(config.vocab_size)
296 - self.make_empty_intermediate_tensors = (
297 - self.gpt_neox.make_empty_intermediate_tensors
298 - )
299 -
300 - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
301 - return self.gpt_neox.embed_input_ids(input_ids)
302 -
303 - def forward(
304 - self,
305 - input_ids: torch.Tensor | None,
306 - positions: torch.Tensor,
307 - intermediate_tensors: IntermediateTensors | None = None,
308 - inputs_embeds: torch.Tensor | None = None,
309 - ) -> torch.Tensor | IntermediateTensors:
310 - hidden_states = self.gpt_neox(
311 - input_ids, positions, intermediate_tensors, inputs_embeds
312 - )
313 - return hidden_states
314 -
315 - def compute_logits(
316 - self,
317 - hidden_states: torch.Tensor,
318 - ) -> torch.Tensor | None:
319 - logits = self.logits_processor(self.embed_out, hidden_states)
320 - return logits
321 -
322 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
323 - loader = AutoWeightsLoader(self)
324 - return loader.load_weights(weights)
vllm/model_executor/models/jais2.py deleted
-441
@@ -1,441 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -# Adapted from
5 -# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/llama/modeling_llama.py
6 -# Copyright 2023 The vLLM team.
7 -# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
8 -#
9 -# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
10 -# and OPT implementations in this library. It has been modified from its
11 -# original forms to accommodate minor architectural differences compared
12 -# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
13 -#
14 -# Licensed under the Apache License, Version 2.0 (the "License");
15 -# you may not use this file except in compliance with the License.
16 -# You may obtain a copy of the License at
17 -#
18 -# http://www.apache.org/licenses/LICENSE-2.0
19 -#
20 -# Unless required by applicable law or agreed to in writing, software
21 -# distributed under the License is distributed on an "AS IS" BASIS,
22 -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
23 -# See the License for the specific language governing permissions and
24 -# limitations under the License.
25 -
26 -"""Inference-only Jais2 model compatible with HuggingFace weights."""
27 -
28 -from collections.abc import Iterable
29 -
30 -import torch
31 -from torch import nn
32 -from transformers import Jais2Config
33 -
34 -from vllm.compilation.decorators import support_torch_compile
35 -from vllm.config import CacheConfig, VllmConfig
36 -from vllm.distributed import (
37 - get_pp_group,
38 - get_tensor_model_parallel_world_size,
39 -)
40 -from vllm.model_executor.layers.activation import ReLUSquaredActivation
41 -from vllm.model_executor.layers.attention import Attention
42 -from vllm.model_executor.layers.linear import (
43 - ColumnParallelLinear,
44 - QKVParallelLinear,
45 - RowParallelLinear,
46 -)
47 -from vllm.model_executor.layers.logits_processor import LogitsProcessor
48 -from vllm.model_executor.layers.quantization import QuantizationConfig
49 -from vllm.model_executor.layers.rotary_embedding import get_rope
50 -from vllm.model_executor.layers.vocab_parallel_embedding import (
51 - ParallelLMHead,
52 - VocabParallelEmbedding,
53 -)
54 -from vllm.sequence import IntermediateTensors
55 -
56 -from .interfaces import SupportsLoRA, SupportsPP
57 -from .utils import (
58 - AutoWeightsLoader,
59 - PPMissingLayer,
60 - WeightsMapper,
61 - extract_layer_index,
62 - make_empty_intermediate_tensors_factory,
63 - make_layers,
64 - maybe_prefix,
65 -)
66 -
67 -
68 -class Jais2MLP(nn.Module):
69 - def __init__(
70 - self,
71 - hidden_size: int,
72 - intermediate_size: int,
73 - hidden_act: str,
74 - quant_config: QuantizationConfig | None = None,
75 - bias: bool = False,
76 - prefix: str = "",
77 - ) -> None:
78 - super().__init__()
79 - self.up_proj = ColumnParallelLinear(
80 - input_size=hidden_size,
81 - output_size=intermediate_size,
82 - bias=bias,
83 - quant_config=quant_config,
84 - prefix=f"{prefix}.up_proj",
85 - )
86 - self.down_proj = RowParallelLinear(
87 - input_size=intermediate_size,
88 - output_size=hidden_size,
89 - bias=bias,
90 - quant_config=quant_config,
91 - prefix=f"{prefix}.down_proj",
92 - )
93 - self.act_fn = ReLUSquaredActivation()
94 -
95 - def forward(self, x):
96 - x, _ = self.up_proj(x)
97 - x = self.act_fn(x)
98 - x, _ = self.down_proj(x)
99 - return x
100 -
101 -
102 -class Jais2Attention(nn.Module):
103 - def __init__(
104 - self,
105 - config: Jais2Config,
106 - hidden_size: int,
107 - num_heads: int,
108 - num_kv_heads: int,
109 - max_position_embeddings: int = 8192,
110 - quant_config: QuantizationConfig | None = None,
111 - bias: bool = False,
112 - cache_config: CacheConfig | None = None,
113 - prefix: str = "",
114 - ) -> None:
115 - super().__init__()
116 - layer_idx = extract_layer_index(prefix)
117 - self.hidden_size = hidden_size
118 - tp_size = get_tensor_model_parallel_world_size()
119 - self.total_num_heads = num_heads
120 - assert self.total_num_heads % tp_size == 0
121 - self.num_heads = self.total_num_heads // tp_size
122 - self.total_num_kv_heads = num_kv_heads
123 - if self.total_num_kv_heads >= tp_size:
124 - # Number of KV heads is greater than TP size, so we partition
125 - # the KV heads across multiple tensor parallel GPUs.
126 - assert self.total_num_kv_heads % tp_size == 0
127 - else:
128 - # Number of KV heads is less than TP size, so we replicate
129 - # the KV heads across multiple tensor parallel GPUs.
130 - assert tp_size % self.total_num_kv_heads == 0
131 - self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
132 - # MistralConfig has an optional head_dim introduced by Mistral-Nemo
133 - self.head_dim = getattr(
134 - config, "head_dim", self.hidden_size // self.total_num_heads
135 - )
136 - self.q_size = self.num_heads * self.head_dim
137 - self.kv_size = self.num_kv_heads * self.head_dim
138 - self.scaling = self.head_dim**-0.5
139 - self.max_position_embeddings = max_position_embeddings
140 -
141 - self.qkv_proj = QKVParallelLinear(
142 - hidden_size=hidden_size,
143 - head_size=self.head_dim,
144 - total_num_heads=self.total_num_heads,
145 - total_num_kv_heads=self.total_num_kv_heads,
146 - bias=bias,
147 - quant_config=quant_config,
148 - prefix=f"{prefix}.qkv_proj",
149 - )
150 -
151 - self.o_proj = RowParallelLinear(
152 - input_size=self.total_num_heads * self.head_dim,
153 - output_size=hidden_size,
154 - bias=bias,
155 - quant_config=quant_config,
156 - prefix=f"{prefix}.o_proj",
157 - )
158 -
159 - is_neox_style = True
160 - self.rotary_emb = get_rope(
161 - self.head_dim,
162 - max_position=max_position_embeddings,
163 - rope_parameters=getattr(config, "rope_parameters", None),
164 - is_neox_style=is_neox_style,
165 - )
166 -
167 - if hasattr(config, "interleaved_sliding_window"):
168 - interleaved_sliding_window = config.interleaved_sliding_window
169 - if isinstance(interleaved_sliding_window, int):
170 - sliding_window = interleaved_sliding_window
171 - elif isinstance(interleaved_sliding_window, list):
172 - sw_idx = layer_idx % len(interleaved_sliding_window)
173 - sliding_window = interleaved_sliding_window[sw_idx]
174 - else:
175 - raise ValueError(
176 - f"{type(interleaved_sliding_window)} is not supported."
177 - )
178 - else:
179 - sliding_window = None
180 -
181 - self.attn = Attention(
182 - self.num_heads,
183 - self.head_dim,
184 - self.scaling,
185 - num_kv_heads=self.num_kv_heads,
186 - cache_config=cache_config,
187 - quant_config=quant_config,
188 - per_layer_sliding_window=sliding_window,
189 - prefix=f"{prefix}.attn",
190 - )
191 -
192 - def forward(
193 - self,
194 - positions: torch.Tensor,
195 - hidden_states: torch.Tensor,
196 - ) -> torch.Tensor:
197 - qkv, _ = self.qkv_proj(hidden_states)
198 - q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
199 - q, k = self.rotary_emb(positions, q, k)
200 - attn_output = self.attn(q, k, v)
201 - output, _ = self.o_proj(attn_output)
202 - return output
203 -
204 -
205 -class Jais2DecoderLayer(nn.Module):
206 - def __init__(
207 - self,
208 - vllm_config: VllmConfig,
209 - config: Jais2Config,
210 - prefix: str = "",
211 - ) -> None:
212 - super().__init__()
213 -
214 - config = config or vllm_config.model_config.hf_config
215 - cache_config = vllm_config.cache_config
216 - quant_config = self.get_quant_config(vllm_config)
217 -
218 - self.hidden_size = config.hidden_size
219 - max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
220 - # Support abacusai/Smaug-72B-v0.1 with attention_bias
221 - attention_bias = getattr(config, "attention_bias", False) or getattr(
222 - config, "bias", False
223 - )
224 - self.self_attn = Jais2Attention(
225 - config=config,
226 - hidden_size=self.hidden_size,
227 - num_heads=config.num_attention_heads,
228 - num_kv_heads=getattr(
229 - config, "num_key_value_heads", config.num_attention_heads
230 - ),
231 - max_position_embeddings=max_position_embeddings,
232 - quant_config=quant_config,
233 - bias=attention_bias,
234 - cache_config=cache_config,
235 - prefix=f"{prefix}.self_attn",
236 - )
237 - self.mlp = Jais2MLP(
238 - hidden_size=self.hidden_size,
239 - intermediate_size=config.intermediate_size,
240 - hidden_act=config.hidden_act,
241 - quant_config=quant_config,
242 - bias=getattr(config, "mlp_bias", False),
243 - prefix=f"{prefix}.mlp",
244 - )
245 - self.input_layernorm = nn.LayerNorm(
246 - config.hidden_size, eps=config.layer_norm_eps
247 - )
248 - self.post_attention_layernorm = nn.LayerNorm(
249 - config.hidden_size, eps=config.layer_norm_eps
250 - )
251 -
252 - def forward(
253 - self,
254 - positions: torch.Tensor,
255 - hidden_states: torch.Tensor,
256 - residual: torch.Tensor | None,
257 - ) -> tuple[torch.Tensor, torch.Tensor]:
258 - # Self Attention
259 - if residual is None:
260 - residual = hidden_states
261 - hidden_states = self.input_layernorm(hidden_states)
262 - else:
263 - hidden_states, residual = (
264 - self.input_layernorm(hidden_states + residual),
265 - hidden_states + residual,
266 - )
267 - hidden_states = self.self_attn(
268 - positions=positions,
269 - hidden_states=hidden_states,
270 - )
271 -
272 - # Fully Connected
273 - hidden_states, residual = (
274 - self.post_attention_layernorm(hidden_states + residual),
275 - hidden_states + residual,
276 - )
277 - hidden_states = self.mlp(hidden_states)
278 - return hidden_states, residual
279 -
280 - def get_quant_config(self, vllm_config: VllmConfig) -> QuantizationConfig | None:
281 - """Get quantization config for this layer. Override in subclasses."""
282 - return vllm_config.quant_config
283 -
284 -
285 -@support_torch_compile
286 -class Jais2Model(nn.Module):
287 - def __init__(
288 - self,
289 - vllm_config: VllmConfig,
290 - prefix: str = "",
291 - layer_type: type[nn.Module] = Jais2DecoderLayer,
292 - ):
293 - super().__init__()
294 -
295 - config = vllm_config.model_config.hf_config
296 - quant_config = vllm_config.quant_config
297 -
298 - self.config = config
299 - self.quant_config = quant_config
300 -
301 - self.vocab_size = config.vocab_size
302 - self.org_vocab_size = config.vocab_size
303 - if get_pp_group().is_first_rank or (
304 - config.tie_word_embeddings and get_pp_group().is_last_rank
305 - ):
306 - self.embed_tokens = VocabParallelEmbedding(
307 - self.vocab_size,
308 - config.hidden_size,
309 - org_num_embeddings=config.vocab_size,
310 - quant_config=quant_config,
311 - )
312 - else:
313 - self.embed_tokens = PPMissingLayer()
314 - self.start_layer, self.end_layer, self.layers = make_layers(
315 - config.num_hidden_layers,
316 - lambda prefix: layer_type(
317 - config=config,
318 - vllm_config=vllm_config,
319 - prefix=prefix,
320 - ),
321 - prefix=f"{prefix}.layers",
322 - )
323 - if get_pp_group().is_last_rank:
324 - self.norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
325 - else:
326 - self.norm = PPMissingLayer()
327 -
328 - self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
329 - ["hidden_states", "residual"], config.hidden_size
330 - )
331 -
332 - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
333 - return self.embed_tokens(input_ids)
334 -
335 - def forward(
336 - self,
337 - input_ids: torch.Tensor | None,
338 - positions: torch.Tensor,
339 - intermediate_tensors: IntermediateTensors | None,
340 - inputs_embeds: torch.Tensor | None = None,
341 - ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
342 - if get_pp_group().is_first_rank:
343 - if inputs_embeds is not None:
344 - hidden_states = inputs_embeds
345 - else:
346 - hidden_states = self.embed_input_ids(input_ids)
347 - residual = None
348 - else:
349 - assert intermediate_tensors is not None
350 - hidden_states = intermediate_tensors["hidden_states"]
351 - residual = intermediate_tensors["residual"]
352 -
353 - for i in range(self.start_layer, self.end_layer):
354 - layer = self.layers[i]
355 - hidden_states, residual = layer(positions, hidden_states, residual)
356 -
357 - if not get_pp_group().is_last_rank:
358 - return IntermediateTensors(
359 - {"hidden_states": hidden_states, "residual": residual}
360 - )
361 -
362 - hidden_states, _ = self.norm(hidden_states + residual), residual
363 - return hidden_states
364 -
365 -
366 -class Jais2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
367 - hf_to_vllm_mapper = WeightsMapper(
368 - orig_to_new_stacked={
369 - # weight_name: (param_name, shard_id)
370 - ".q_proj": (".qkv_proj", "q"),
371 - ".k_proj": (".qkv_proj", "k"),
372 - ".v_proj": (".qkv_proj", "v"),
373 - }
374 - )
375 - packed_modules_mapping = {
376 - "qkv_proj": ["q_proj", "k_proj", "v_proj"],
377 - }
378 -
379 - embedding_modules = {
380 - "embed_tokens": "input_embeddings",
381 - "lm_head": "output_embeddings",
382 - }
383 -
384 - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
385 - super().__init__()
386 - config = vllm_config.model_config.hf_config
387 - quant_config = vllm_config.quant_config
388 - self.config = config
389 - self.model = self._init_model(
390 - vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
391 - )
392 -
393 - if get_pp_group().is_last_rank:
394 - self.lm_head = ParallelLMHead(
395 - config.vocab_size,
396 - config.hidden_size,
397 - quant_config=quant_config,
398 - prefix=maybe_prefix(prefix, "lm_head"),
399 - )
400 - if config.tie_word_embeddings:
401 - self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)
402 -
403 - logit_scale = getattr(config, "logit_scale", 1.0)
404 - self.logits_processor = LogitsProcessor(
405 - config.vocab_size, scale=logit_scale
406 - )
407 - else:
408 - self.lm_head = PPMissingLayer()
409 -
410 - self.make_empty_intermediate_tensors = (
411 - self.model.make_empty_intermediate_tensors
412 - )
413 -
414 - def _init_model(self, vllm_config: VllmConfig, prefix: str = ""):
415 - return Jais2Model(vllm_config=vllm_config, prefix=prefix)
416 -
417 - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
418 - return self.model.embed_input_ids(input_ids)
419 -
420 - def forward(
421 - self,
422 - input_ids: torch.Tensor | None,
423 - positions: torch.Tensor,
424 - intermediate_tensors: IntermediateTensors | None = None,
425 - inputs_embeds: torch.Tensor | None = None,
426 - ) -> torch.Tensor | IntermediateTensors:
427 - model_output = self.model(
428 - input_ids, positions, intermediate_tensors, inputs_embeds
429 - )
430 - return model_output
431 -
432 - def compute_logits(
433 - self,
434 - hidden_states: torch.Tensor,
435 - ) -> torch.Tensor | None:
436 - logits = self.logits_processor(self.lm_head, hidden_states)
437 - return logits
438 -
439 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
440 - loader = AutoWeightsLoader(self)
441 - return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/phi.py deleted
-325
@@ -1,325 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -# Adapted from
5 -# https://huggingface.co/microsoft/phi-1_5/blob/main/modeling_phi.py
6 -# Copyright 2023 The vLLM team.
7 -# Copyright (c) Microsoft Corporation.
8 -# Licensed under the MIT license.
9 -#
10 -# BSD 3-Clause License
11 -#
12 -# Copyright (c) 2022, Tri Dao, trid@cs.stanford.edu.
13 -# All rights reserved.
14 -#
15 -# Redistribution and use in source and binary forms, with or without
16 -# modification, are permitted provided that the following conditions are met:
17 -#
18 -# * Redistributions of source code must retain the above copyright notice, this
19 -# list of conditions and the following disclaimer.
20 -#
21 -# * Redistributions in binary form must reproduce the above copyright notice,
22 -# this list of conditions and the following disclaimer in the documentation
23 -# and/or other materials provided with the distribution.
24 -#
25 -# * Neither the name of the copyright holder nor the names of its
26 -# contributors may be used to endorse or promote products derived from
27 -# this software without specific prior written permission.
28 -#
29 -# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
30 -# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
31 -# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
32 -# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
33 -# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
34 -# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
35 -# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
36 -# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
37 -# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
38 -# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
39 -"""Inference-only Phi-1.5 model compatible with HuggingFace weights."""
40 -
41 -from collections.abc import Iterable
42 -from itertools import islice
43 -
44 -import torch
45 -from torch import nn
46 -from transformers import PhiConfig
47 -
48 -from vllm.compilation.decorators import support_torch_compile
49 -from vllm.config import CacheConfig, VllmConfig
50 -from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
51 -from vllm.model_executor.layers.activation import get_act_fn
52 -from vllm.model_executor.layers.attention import Attention
53 -from vllm.model_executor.layers.linear import (
54 - ColumnParallelLinear,
55 - QKVParallelLinear,
56 - RowParallelLinear,
57 -)
58 -from vllm.model_executor.layers.logits_processor import LogitsProcessor
59 -from vllm.model_executor.layers.quantization import QuantizationConfig
60 -from vllm.model_executor.layers.rotary_embedding import get_rope
61 -from vllm.model_executor.layers.vocab_parallel_embedding import (
62 - ParallelLMHead,
63 - VocabParallelEmbedding,
64 -)
65 -from vllm.sequence import IntermediateTensors
66 -
67 -from .interfaces import SupportsLoRA, SupportsPP
68 -from .utils import (
69 - AutoWeightsLoader,
70 - WeightsMapper,
71 - make_empty_intermediate_tensors_factory,
72 - make_layers,
73 - maybe_prefix,
74 -)
75 -
76 -
77 -class PhiAttention(nn.Module):
78 - def __init__(
79 - self,
80 - config: PhiConfig,
81 - cache_config: CacheConfig | None = None,
82 - quant_config: QuantizationConfig | None = None,
83 - prefix: str = "",
84 - ):
85 - super().__init__()
86 - self.hidden_size = config.hidden_size
87 - self.head_size = self.hidden_size // config.num_attention_heads
88 -
89 - tensor_model_parallel_world_size = get_tensor_model_parallel_world_size()
90 - assert config.num_attention_heads % tensor_model_parallel_world_size == 0
91 - self.num_heads = config.num_attention_heads // tensor_model_parallel_world_size
92 -
93 - # pylint: disable=C0103
94 - self.qkv_proj = QKVParallelLinear(
95 - self.hidden_size,
96 - self.head_size,
97 - config.num_attention_heads,
98 - bias=True,
99 - quant_config=quant_config,
100 - prefix=f"{prefix}.qkv_proj",
101 - )
102 - self.dense = RowParallelLinear(
103 - self.hidden_size,
104 - self.hidden_size,
105 - quant_config=quant_config,
106 - prefix=f"{prefix}.dense",
107 - )
108 -
109 - scaling = self.head_size**-0.5
110 -
111 - max_position_embeddings = getattr(config, "max_position_embeddings", 2048)
112 - self.rotary_emb = get_rope(
113 - self.head_size,
114 - max_position=max_position_embeddings,
115 - rope_parameters=config.rope_parameters,
116 - )
117 - self.attn = Attention(
118 - self.num_heads,
119 - self.head_size,
120 - scaling,
121 - cache_config=cache_config,
122 - quant_config=quant_config,
123 - prefix=f"{prefix}.attn",
124 - )
125 -
126 - def forward(
127 - self,
128 - position_ids: torch.Tensor,
129 - hidden_states: torch.Tensor,
130 - ) -> torch.Tensor:
131 - qkv, _ = self.qkv_proj(hidden_states)
132 - q, k, v = qkv.chunk(chunks=3, dim=-1)
133 - q, k = self.rotary_emb(position_ids, q, k)
134 - attn_output = self.attn(q, k, v)
135 - output, _ = self.dense(attn_output)
136 - return output
137 -
138 -
139 -class PhiMLP(nn.Module):
140 - def __init__(
141 - self,
142 - config: PhiConfig,
143 - quant_config: QuantizationConfig | None = None,
144 - prefix: str = "",
145 - ):
146 - super().__init__()
147 -
148 - n_inner = getattr(config, "n_inner", None)
149 - n_inner = n_inner if n_inner is not None else 4 * config.hidden_size
150 -
151 - self.fc1 = ColumnParallelLinear(
152 - config.hidden_size,
153 - n_inner,
154 - quant_config=quant_config,
155 - prefix=f"{prefix}.fc1",
156 - )
157 - self.fc2 = RowParallelLinear(
158 - n_inner,
159 - config.hidden_size,
160 - quant_config=quant_config,
161 - prefix=f"{prefix}.fc2",
162 - )
163 - self.act = get_act_fn(config.hidden_act)
164 -
165 - def forward(self, hidden_states):
166 - hidden_states, _ = self.fc1(hidden_states)
167 - hidden_states = self.act(hidden_states)
168 - hidden_states, _ = self.fc2(hidden_states)
169 - return hidden_states
170 -
171 -
172 -class PhiLayer(nn.Module):
173 - def __init__(
174 - self,
175 - config: PhiConfig,
176 - cache_config: CacheConfig | None = None,
177 - quant_config: QuantizationConfig | None = None,
178 - prefix: str = "",
179 - ):
180 - super().__init__()
181 - self.input_layernorm = nn.LayerNorm(
182 - config.hidden_size, eps=config.layer_norm_eps
183 - )
184 - self.self_attn = PhiAttention(
185 - config, cache_config, quant_config, prefix=f"{prefix}.self_attn"
186 - )
187 - self.mlp = PhiMLP(config, quant_config, prefix=f"{prefix}.mlp")
188 -
189 - def forward(
190 - self,
191 - position_ids: torch.Tensor,
192 - hidden_states: torch.Tensor,
193 - ) -> torch.Tensor:
194 - residual = hidden_states
195 - hidden_states = self.input_layernorm(hidden_states)
196 - attn_outputs = self.self_attn(
197 - position_ids=position_ids,
198 - hidden_states=hidden_states,
199 - )
200 - feed_forward_hidden_states = self.mlp(hidden_states)
201 - hidden_states = attn_outputs + feed_forward_hidden_states + residual
202 - return hidden_states
203 -
204 -
205 -@support_torch_compile
206 -class PhiModel(nn.Module):
207 - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
208 - super().__init__()
209 -
210 - config = vllm_config.model_config.hf_config
211 - cache_config = vllm_config.cache_config
212 - quant_config = vllm_config.quant_config
213 -
214 - self.config = config
215 - self.quant_config = quant_config
216 - self.embed_tokens = VocabParallelEmbedding(
217 - config.vocab_size, config.hidden_size
218 - )
219 - self.start_layer, self.end_layer, self.layers = make_layers(
220 - config.num_hidden_layers,
221 - lambda prefix: PhiLayer(config, cache_config, quant_config, prefix=prefix),
222 - prefix=f"{prefix}.layers",
223 - )
224 - self.final_layernorm = nn.LayerNorm(
225 - config.hidden_size, eps=config.layer_norm_eps
226 - )
227 - self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
228 - ["hidden_states"], config.hidden_size
229 - )
230 -
231 - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
232 - return self.embed_tokens(input_ids)
233 -
234 - def forward(
235 - self,
236 - input_ids: torch.Tensor | None,
237 - positions: torch.Tensor,
238 - intermediate_tensors: IntermediateTensors | None,
239 - inputs_embeds: torch.Tensor | None = None,
240 - ) -> torch.Tensor | IntermediateTensors:
241 - if get_pp_group().is_first_rank:
242 - if inputs_embeds is not None:
243 - hidden_states = inputs_embeds
244 - else:
245 - hidden_states = self.embed_input_ids(input_ids)
246 - else:
247 - assert intermediate_tensors is not None
248 - hidden_states = intermediate_tensors["hidden_states"]
249 - for layer in islice(self.layers, self.start_layer, self.end_layer):
250 - hidden_states = layer(positions, hidden_states)
251 -
252 - if not get_pp_group().is_last_rank:
253 - return IntermediateTensors({"hidden_states": hidden_states})
254 -
255 - hidden_states = self.final_layernorm(hidden_states)
256 -
257 - return hidden_states
258 -
259 -
260 -class PhiForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
261 - hf_to_vllm_mapper = WeightsMapper(
262 - orig_to_new_stacked={
263 - # weight_name: (param_name, shard_id)
264 - ".q_proj": (".qkv_proj", "q"),
265 - ".k_proj": (".qkv_proj", "k"),
266 - ".v_proj": (".qkv_proj", "v"),
267 - }
268 - )
269 - packed_modules_mapping = {
270 - "qkv_proj": ["q_proj", "k_proj", "v_proj"],
271 - }
272 -
273 - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
274 - super().__init__()
275 - config = vllm_config.model_config.hf_config
276 - quant_config = vllm_config.quant_config
277 -
278 - self.config = config
279 - # lm_head use bias, cannot share word embeddings
280 - assert not config.tie_word_embeddings
281 -
282 - self.quant_config = quant_config
283 -
284 - self.model = PhiModel(
285 - vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
286 - )
287 -
288 - self.lm_head = ParallelLMHead(
289 - config.vocab_size,
290 - config.hidden_size,
291 - bias=True,
292 - quant_config=quant_config,
293 - prefix=maybe_prefix(prefix, "lm_head"),
294 - )
295 - self.logits_processor = LogitsProcessor(config.vocab_size)
296 - self.make_empty_intermediate_tensors = (
297 - self.model.make_empty_intermediate_tensors
298 - )
299 -
300 - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
301 - return self.model.embed_input_ids(input_ids)
302 -
303 - def forward(
304 - self,
305 - input_ids: torch.Tensor | None,
306 - positions: torch.Tensor,
307 - intermediate_tensors: IntermediateTensors | None = None,
308 - inputs_embeds: torch.Tensor | None = None,
309 - ) -> torch.Tensor | IntermediateTensors:
310 - hidden_states = self.model(
311 - input_ids, positions, intermediate_tensors, inputs_embeds
312 - )
313 -
314 - return hidden_states
315 -
316 - def compute_logits(
317 - self,
318 - hidden_states: torch.Tensor,
319 - ) -> torch.Tensor | None:
320 - logits = self.logits_processor(self.lm_head, hidden_states, self.lm_head.bias)
321 - return logits
322 -
323 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
324 - loader = AutoWeightsLoader(self)
325 - return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/registry.py
+4 -4
@@ -120,7 +120,6 @@ _TEXT_GENERATION_MODELS = {
120 "GptOssForCausalLM": ("gpt_oss", "GptOssForCausalLM"),
121 "GPT2LMHeadModel": ("gpt2", "GPT2LMHeadModel"),
122 "GPTJForCausalLM": ("gpt_j", "GPTJForCausalLM"),
123 - "GPTNeoXForCausalLM": ("gpt_neox", "GPTNeoXForCausalLM"),
123 "GraniteForCausalLM": ("granite", "GraniteForCausalLM"),
124 "GraniteMoeForCausalLM": ("granitemoe", "GraniteMoeForCausalLM"),
125 "GraniteMoeHybridForCausalLM": ("granitemoehybrid", "GraniteMoeHybridForCausalLM"),
@@ -136,7 +135,6 @@ _TEXT_GENERATION_MODELS = {
135 "InternLM3ForCausalLM": ("llama", "LlamaForCausalLM"),
136 "IQuestCoderForCausalLM": ("llama", "LlamaForCausalLM"),
137 "IQuestLoopCoderForCausalLM": ("iquest_loopcoder", "IQuestLoopCoderForCausalLM"),
139 - "Jais2ForCausalLM": ("jais2", "Jais2ForCausalLM"),
138 "JambaForCausalLM": ("jamba", "JambaForCausalLM"),
139 "K2HorizonForCausalLM": ("k2_horizon", "K2HorizonForCausalLM"),
140 "KimiLinearForCausalLM": (
@@ -188,7 +186,6 @@ _TEXT_GENERATION_MODELS = {
186 "PanguProMoEV2ForCausalLM": ("openpangu", "PanguProMoEV2ForCausalLM"),
187 "PanguUltraMoEForCausalLM": ("openpangu", "PanguUltraMoEForCausalLM"),
188 "Param2MoEForCausalLM": ("param2moe", "Param2MoEForCausalLM"),
191 - "PhiForCausalLM": ("phi", "PhiForCausalLM"),
189 "Phi3ForCausalLM": ("phi3", "Phi3ForCausalLM"),
190 "PhiMoEForCausalLM": ("phimoe", "PhiMoEForCausalLM"),
191 "Plamo3ForCausalLM": ("plamo3", "Plamo3ForCausalLM"),
@@ -200,7 +197,6 @@ _TEXT_GENERATION_MODELS = {
197 "Qwen3_5MoeForCausalLM": ("qwen3_5", "Qwen3_5MoeForCausalLM"),
198 "SarvamMoEForCausalLM": ("sarvam", "SarvamMoEForCausalLM"),
199 "SarvamMLAForCausalLM": ("sarvam", "SarvamMLAForCausalLM"),
203 - "SeedOssForCausalLM": ("seed_oss", "SeedOssForCausalLM"),
200 "Step1ForCausalLM": ("step1", "Step1ForCausalLM"),
201 "Step3TextForCausalLM": ("step3_text", "Step3TextForCausalLM"),
202 "Step3p5ForCausalLM": ("step3p5", "Step3p5ForCausalLM"),
@@ -709,13 +705,17 @@ _TRANSFORMERS_SUPPORTED_MODELS = {
705 "FlexOlmoForCausalLM": ("transformers", "TransformersMoEForCausalLM"),
706 "GlmForCausalLM": ("transformers", "TransformersForCausalLM"),
707 "GPTBigCodeForCausalLM": ("transformers", "TransformersForCausalLM"),
708 + "GPTNeoXForCausalLM": ("transformers", "TransformersForCausalLM"),
709 "HunYuanDenseV1ForCausalLM": ("transformers", "TransformersForCausalLM"),
710 "HunYuanMoEV1ForCausalLM": ("transformers", "TransformersMoEForCausalLM"),
711 + "Jais2ForCausalLM": ("transformers", "TransformersForCausalLM"),
712 "MellumForCausalLM": ("transformers", "TransformersMoEForCausalLM"),
713 "NanbeigeForCausalLM": ("transformers", "TransformersForCausalLM"),
714 "OlmoForCausalLM": ("transformers", "TransformersForCausalLM"),
715 "Olmo2ForCausalLM": ("transformers", "TransformersForCausalLM"),
716 "Olmo3ForCausalLM": ("transformers", "TransformersForCausalLM"),
717 + "PhiForCausalLM": ("transformers", "TransformersForCausalLM"),
718 + "SeedOssForCausalLM": ("transformers", "TransformersForCausalLM"),
719 "SmolLM3ForCausalLM": ("transformers", "TransformersForCausalLM"),
720 "Starcoder2ForCausalLM": ("transformers", "TransformersForCausalLM"),
721 "VaultGemmaForCausalLM": ("transformers", "TransformersForCausalLM"),
vllm/model_executor/models/seed_oss.py deleted
-432
@@ -1,432 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -# Copyright 2025 The Seed team.
5 -# Copyright 2023 The vLLM team.
6 -# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
7 -#
8 -# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
9 -# and OPT implementations in this library. It has been modified from its
10 -# original forms to accommodate minor architectural differences compared
11 -# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
12 -#
13 -# Licensed under the Apache License, Version 2.0 (the "License");
14 -# you may not use this file except in compliance with the License.
15 -# You may obtain a copy of the License at
16 -#
17 -# http://www.apache.org/licenses/LICENSE-2.0
18 -#
19 -# Unless required by applicable law or agreed to in writing, software
20 -# distributed under the License is distributed on an "AS IS" BASIS,
21 -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
22 -# See the License for the specific language governing permissions and
23 -# limitations under the License.
24 -"""Inference-only SeedOss model compatible with HuggingFace weights."""
25 -
26 -from collections.abc import Iterable
27 -from itertools import islice
28 -
29 -import torch
30 -from torch import nn
31 -from transformers import PreTrainedConfig as SeedOssConfig
32 -
33 -from vllm.compilation.decorators import support_torch_compile
34 -from vllm.config import CacheConfig, VllmConfig
35 -from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
36 -from vllm.logger import init_logger
37 -from vllm.model_executor.layers.activation import SiluAndMul
38 -from vllm.model_executor.layers.attention import Attention
39 -from vllm.model_executor.layers.layernorm import RMSNorm
40 -from vllm.model_executor.layers.linear import (
41 - MergedColumnParallelLinear,
42 - QKVParallelLinear,
43 - RowParallelLinear,
44 -)
45 -from vllm.model_executor.layers.logits_processor import LogitsProcessor
46 -from vllm.model_executor.layers.quantization import QuantizationConfig
47 -from vllm.model_executor.layers.rotary_embedding import get_rope
48 -from vllm.model_executor.layers.vocab_parallel_embedding import (
49 - ParallelLMHead,
50 - VocabParallelEmbedding,
51 -)
52 -from vllm.sequence import IntermediateTensors
53 -from vllm.transformers_utils.config import set_default_rope_theta
54 -from vllm.v1.attention.backend import AttentionType
55 -
56 -from .interfaces import SupportsLoRA, SupportsPP
57 -from .utils import (
58 - AutoWeightsLoader,
59 - PPMissingLayer,
60 - WeightsMapper,
61 - make_empty_intermediate_tensors_factory,
62 - make_layers,
63 - maybe_prefix,
64 -)
65 -
66 -logger = init_logger(__name__)
67 -
68 -
69 -class SeedOssMLP(nn.Module):
70 - def __init__(
71 - self,
72 - hidden_size: int,
73 - intermediate_size: int,
74 - hidden_act: str,
75 - quant_config: QuantizationConfig | None = None,
76 - prefix: str = "",
77 - ) -> None:
78 - super().__init__()
79 - self.gate_up_proj = MergedColumnParallelLinear(
80 - hidden_size,
81 - [intermediate_size] * 2,
82 - bias=False,
83 - quant_config=quant_config,
84 - prefix=f"{prefix}.gate_up_proj",
85 - )
86 - self.down_proj = RowParallelLinear(
87 - intermediate_size,
88 - hidden_size,
89 - bias=False,
90 - quant_config=quant_config,
91 - prefix=f"{prefix}.down_proj",
92 - )
93 - if hidden_act != "silu":
94 - raise ValueError(
95 - f"Unsupported activation: {hidden_act}. Only silu is supported for now."
96 - )
97 - self.act_fn = SiluAndMul()
98 -
99 - def forward(self, x):
100 - gate_up, _ = self.gate_up_proj(x)
101 - x = self.act_fn(gate_up)
102 - x, _ = self.down_proj(x)
103 - return x
104 -
105 -
106 -class SeedOssAttention(nn.Module):
107 - def __init__(
108 - self,
109 - hidden_size: int,
110 - num_heads: int,
111 - num_kv_heads: int,
112 - head_dim: int,
113 - rope_parameters: dict,
114 - max_position: int = 4096 * 32,
115 - cache_config: CacheConfig | None = None,
116 - quant_config: QuantizationConfig | None = None,
117 - prefix: str = "",
118 - attn_type: str = AttentionType.DECODER,
119 - ) -> None:
120 - super().__init__()
121 - self.hidden_size = hidden_size
122 - tp_size = get_tensor_model_parallel_world_size()
123 - self.total_num_heads = num_heads
124 - assert self.total_num_heads % tp_size == 0
125 - self.num_heads = self.total_num_heads // tp_size
126 - self.total_num_kv_heads = num_kv_heads
127 - self.head_dim = head_dim
128 - if self.total_num_kv_heads >= tp_size:
129 - # Number of KV heads is greater than TP size, so we partition
130 - # the KV heads across multiple tensor parallel GPUs.
131 - assert self.total_num_kv_heads % tp_size == 0
132 - else:
133 - # Number of KV heads is less than TP size, so we replicate
134 - # the KV heads across multiple tensor parallel GPUs.
135 - assert tp_size % self.total_num_kv_heads == 0
136 - self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
137 - self.q_size = self.num_heads * self.head_dim
138 - self.kv_size = self.num_kv_heads * self.head_dim
139 - self.scaling = self.head_dim**-0.5
140 -
141 - self.qkv_proj = QKVParallelLinear(
142 - hidden_size,
143 - self.head_dim,
144 - self.total_num_heads,
145 - self.total_num_kv_heads,
146 - bias=True,
147 - quant_config=quant_config,
148 - prefix=f"{prefix}.qkv_proj",
149 - )
150 - self.o_proj = RowParallelLinear(
151 - self.total_num_heads * self.head_dim,
152 - hidden_size,
153 - bias=False,
154 - quant_config=quant_config,
155 - prefix=f"{prefix}.o_proj",
156 - )
157 -
158 - self.rotary_emb = get_rope(
159 - self.head_dim,
160 - max_position=max_position,
161 - rope_parameters=rope_parameters,
162 - )
163 - self.attn = Attention(
164 - self.num_heads,
165 - self.head_dim,
166 - self.scaling,
167 - num_kv_heads=self.num_kv_heads,
168 - cache_config=cache_config,
169 - quant_config=quant_config,
170 - attn_type=attn_type,
171 - prefix=f"{prefix}.attn",
172 - )
173 -
174 - def forward(
175 - self,
176 - positions: torch.Tensor,
177 - hidden_states: torch.Tensor,
178 - ) -> torch.Tensor:
179 - qkv, _ = self.qkv_proj(hidden_states)
180 - q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
181 - q, k = self.rotary_emb(positions, q, k)
182 - attn_output = self.attn(q, k, v)
183 - output, _ = self.o_proj(attn_output)
184 - return output
185 -
186 -
187 -class SeedOssDecoderLayer(nn.Module):
188 - def __init__(
189 - self,
190 - config: SeedOssConfig,
191 - cache_config: CacheConfig | None = None,
192 - quant_config: QuantizationConfig | None = None,
193 - prefix: str = "",
194 - ) -> None:
195 - super().__init__()
196 - self.hidden_size = config.hidden_size
197 - set_default_rope_theta(config, default_theta=1000000)
198 -
199 - # By default, SeedOss uses causal attention as it is a
200 - # decoder-only model.
201 - # You can override the HF config with `is_causal=False` to enable
202 - # bidirectional attention, which is used in some embedding models
203 - if getattr(config, "is_causal", True):
204 - attn_type = AttentionType.DECODER
205 - else:
206 - attn_type = AttentionType.ENCODER_ONLY
207 -
208 - self.self_attn = SeedOssAttention(
209 - hidden_size=self.hidden_size,
210 - num_heads=config.num_attention_heads,
211 - max_position=config.max_position_embeddings,
212 - num_kv_heads=config.num_key_value_heads,
213 - head_dim=config.head_dim,
214 - cache_config=cache_config,
215 - quant_config=quant_config,
216 - rope_parameters=config.rope_parameters,
217 - prefix=f"{prefix}.self_attn",
218 - attn_type=attn_type,
219 - )
220 - self.mlp = SeedOssMLP(
221 - hidden_size=self.hidden_size,
222 - intermediate_size=config.intermediate_size,
223 - hidden_act=config.hidden_act,
224 - quant_config=quant_config,
225 - prefix=f"{prefix}.mlp",
226 - )
227 - self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
228 - self.post_attention_layernorm = RMSNorm(
229 - config.hidden_size, eps=config.rms_norm_eps
230 - )
231 -
232 - def forward(
233 - self,
234 - positions: torch.Tensor,
235 - hidden_states: torch.Tensor,
236 - residual: torch.Tensor | None,
237 - ) -> tuple[torch.Tensor, torch.Tensor]:
238 - # Self Attention
239 - if residual is None:
240 - residual = hidden_states
241 - hidden_states = self.input_layernorm(hidden_states)
242 - else:
243 - hidden_states, residual = self.input_layernorm(hidden_states, residual)
244 - hidden_states = self.self_attn(
245 - positions=positions,
246 - hidden_states=hidden_states,
247 - )
248 -
249 - # Fully Connected
250 - hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
251 - hidden_states = self.mlp(hidden_states)
252 - return hidden_states, residual
253 -
254 -
255 -@support_torch_compile(
256 - dynamic_arg_dims={
257 - "input_ids": 0,
258 - "positions": -1,
259 - "intermediate_tensors": 0,
260 - "inputs_embeds": 0,
261 - }
262 -)
263 -class SeedOssModel(nn.Module):
264 - def __init__(
265 - self,
266 - *,
267 - vllm_config: VllmConfig,
268 - prefix: str = "",
269 - decoder_layer_type: type[nn.Module] = SeedOssDecoderLayer,
270 - ):
271 - super().__init__()
272 -
273 - config = vllm_config.model_config.hf_config
274 - cache_config = vllm_config.cache_config
275 - quant_config = vllm_config.quant_config
276 -
277 - # TODO (@robertgshaw2): see if this can be moved out
278 - if cache_config.sliding_window is not None and hasattr(
279 - config, "max_window_layers"
280 - ):
281 - assert config.max_window_layers == config.num_hidden_layers, (
282 - "Sliding window for some but all layers is not supported. "
283 - "This model uses sliding window but `max_window_layers` = {} "
284 - "is less than `num_hidden_layers` = {}. Please open an issue "
285 - "to discuss this feature.".format(
286 - config.max_window_layers,
287 - config.num_hidden_layers,
288 - )
289 - )
290 -
291 - self.config = config
292 - self.quant_config = quant_config
293 - self.vocab_size = config.vocab_size
294 -
295 - if get_pp_group().is_first_rank or (
296 - config.tie_word_embeddings and get_pp_group().is_last_rank
297 - ):
298 - self.embed_tokens = VocabParallelEmbedding(
299 - config.vocab_size,
300 - config.hidden_size,
301 - quant_config=quant_config,
302 - prefix=f"{prefix}.embed_tokens",
303 - )
304 - else:
305 - self.embed_tokens = PPMissingLayer()
306 -
307 - # Use the provided decoder layer type or default to SeedDecoderLayer
308 - decoder_layer_type = decoder_layer_type or SeedOssDecoderLayer
309 - self.start_layer, self.end_layer, self.layers = make_layers(
310 - config.num_hidden_layers,
311 - lambda prefix: decoder_layer_type(
312 - config=config,
313 - cache_config=cache_config,
314 - quant_config=quant_config,
315 - prefix=prefix,
316 - ),
317 - prefix=f"{prefix}.layers",
318 - )
319 -
320 - self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
321 - ["hidden_states", "residual"], config.hidden_size
322 - )
323 - if get_pp_group().is_last_rank:
324 - self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
325 - else:
326 - self.norm = PPMissingLayer()
327 -
328 - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
329 - return self.embed_tokens(input_ids)
330 -
331 - def forward(
332 - self,
333 - input_ids: torch.Tensor | None,
334 - positions: torch.Tensor,
335 - intermediate_tensors: IntermediateTensors | None = None,
336 - inputs_embeds: torch.Tensor | None = None,
337 - ) -> torch.Tensor | IntermediateTensors:
338 - if get_pp_group().is_first_rank:
339 - if inputs_embeds is not None:
340 - hidden_states = inputs_embeds
341 - else:
342 - hidden_states = self.embed_input_ids(input_ids)
343 - residual = None
344 - else:
345 - assert intermediate_tensors is not None
346 - hidden_states = intermediate_tensors["hidden_states"]
347 - residual = intermediate_tensors["residual"]
348 - for layer in islice(self.layers, self.start_layer, self.end_layer):
349 - hidden_states, residual = layer(
350 - positions,
351 - hidden_states,
352 - residual,
353 - )
354 - if not get_pp_group().is_last_rank:
355 - return IntermediateTensors(
356 - {"hidden_states": hidden_states, "residual": residual}
357 - )
358 - hidden_states, _ = self.norm(hidden_states, residual)
359 - return hidden_states
360 -
361 -
362 -class SeedOssForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
363 - hf_to_vllm_mapper = WeightsMapper(
364 - orig_to_new_stacked={
365 - # weight_name: (param_name, shard_id)
366 - ".q_proj": (".qkv_proj", "q"),
367 - ".k_proj": (".qkv_proj", "k"),
368 - ".v_proj": (".qkv_proj", "v"),
369 - ".gate_proj": (".gate_up_proj", 0),
370 - ".up_proj": (".gate_up_proj", 1),
371 - }
372 - )
373 - packed_modules_mapping = {
374 - "qkv_proj": ["q_proj", "k_proj", "v_proj"],
375 - "gate_up_proj": ["gate_proj", "up_proj"],
376 - }
377 -
378 - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
379 - super().__init__()
380 - config = vllm_config.model_config.hf_config
381 - quant_config = vllm_config.quant_config
382 -
383 - self.config = config
384 -
385 - self.quant_config = quant_config
386 - self.model = SeedOssModel(
387 - vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
388 - )
389 -
390 - if get_pp_group().is_last_rank:
391 - self.lm_head = ParallelLMHead(
392 - config.vocab_size,
393 - config.hidden_size,
394 - quant_config=quant_config,
395 - prefix=maybe_prefix(prefix, "lm_head"),
396 - )
397 - if config.tie_word_embeddings:
398 - self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)
399 - else:
400 - self.lm_head = PPMissingLayer()
401 -
402 - self.logits_processor = LogitsProcessor(config.vocab_size)
403 -
404 - self.make_empty_intermediate_tensors = (
405 - self.model.make_empty_intermediate_tensors
406 - )
407 -
408 - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
409 - return self.model.embed_input_ids(input_ids)
410 -
411 - def forward(
412 - self,
413 - input_ids: torch.Tensor | None,
414 - positions: torch.Tensor,
415 - intermediate_tensors: IntermediateTensors | None = None,
416 - inputs_embeds: torch.Tensor | None = None,
417 - ) -> torch.Tensor | IntermediateTensors:
418 - hidden_states = self.model(
419 - input_ids, positions, intermediate_tensors, inputs_embeds
420 - )
421 - return hidden_states
422 -
423 - def compute_logits(
424 - self,
425 - hidden_states: torch.Tensor,
426 - ) -> torch.Tensor | None:
427 - logits = self.logits_processor(self.lm_head, hidden_states)
428 - return logits
429 -
430 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
431 - loader = AutoWeightsLoader(self)
432 - return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/transformers/base.py
+9 -1
@@ -32,6 +32,7 @@ from torch import nn
32 from transformers import AutoModel, PreTrainedConfig
33 from transformers.conversion_mapping import (
34 WeightRenaming,
35 + get_checkpoint_conversion_mapping,
36 get_model_conversion_mapping,
37 )
38
@@ -311,7 +312,14 @@ class Base(
312 orig_to_new_renaming: list[WeightRenaming] = []
313 orig_to_new_regex: dict[re.Pattern, str | None] = {}
314
314 - for mapping in get_model_conversion_mapping(self.model):
315 + mappings = get_model_conversion_mapping(self.model)
316 + # Ensure layers that are part of the model, but not AutoModel, get their mapping
317 + instantiated = {type(m).__name__ for m in self.model.modules()}
318 + for arch in self.config.architectures or []:
319 + if arch not in instantiated:
320 + mappings.extend(get_checkpoint_conversion_mapping(arch) or [])
321 +
322 + for mapping in mappings:
323 # Handle weights which have been renamed in Transformers
324 if isinstance(mapping, WeightRenaming):
325 orig_to_new_renaming.append(mapping)
vllm/model_executor/models/transformers/causal.py
+7 -3
@@ -19,6 +19,8 @@
19 from collections.abc import Iterable
20 from typing import TYPE_CHECKING
21
22 +import torch
23 +
24 from vllm.model_executor.layers.logits_processor import LogitsProcessor
25 from vllm.model_executor.layers.vocab_parallel_embedding import (
26 ParallelLMHead,
@@ -30,8 +32,6 @@ from vllm.model_executor.models.utils import PPMissingLayer, maybe_prefix
32 from .base import Base
33
34 if TYPE_CHECKING:
33 - import torch
34 -
35 from vllm.config import VllmConfig
36
37
@@ -68,11 +68,15 @@ class CausalMixin(VllmModelForTextGeneration, Base):
68 def load_weights(self, weights: Iterable[tuple[str, "torch.Tensor"]]) -> set[str]:
69 """A thin wrapper around `Base.load_weights` to handle the lm_head bias."""
70 lm_head_bias = set()
71 + load_device = self.vllm_config.load_config.device or self.device_config.device
72
73 def auto_load_lm_head_bias(weights):
74 for name, weight in weights:
75 if name.endswith("lm_head.bias") and self.pp_group.is_last_rank:
75 - self.lm_head._register_bias()
76 + # load_weights is called outside vLLM's device context,
77 + # so we must enter it again when registering the bias.
78 + with torch.device(load_device):
79 + self.lm_head._register_bias()
80 self.lm_head.bias.weight_loader(self.lm_head.bias, weight)
81 lm_head_bias.add(name)
82 else: