Migrate GPTBigCode and Starcoder2 to the Transformers modeling backend (#30966)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Jul 1, 2026 at 14:41 UTC
a78c15616f927745444ed6a783d98d865643e1ff
5 files changed
+4
-681
docs/models/supported_models.md
+2
-2
@@ -405,7 +405,6 @@ th {
405
| `Glm4MoeLiteForCausalLM` | GLM-4.7-Flash | `zai-org/GLM-4.7-Flash`, etc. | ✅︎ | ✅︎ |
406
| `GlmMoeDsaForCausalLM` | GLM-5, GLM-5.1, GLM-5.2 | `zai-org/GLM-5`, etc. | ✅︎ | ✅︎ |
407
| `GPT2LMHeadModel` | GPT-2 | `openai-community/gpt2`, `openai-community/gpt2-xl`, etc. | | ✅︎ |
408
-| `GPTBigCodeForCausalLM` | StarCoder, SantaCoder, WizardCoder | `bigcode/starcoder`, `bigcode/gpt_bigcode-santacoder`, `WizardLM/WizardCoder-15B-V1.0`, etc. | ✅︎ | ✅︎ |
408
| `GPTJForCausalLM` | GPT-J | `EleutherAI/gpt-j-6b`, `nomic-ai/gpt4all-j`, etc. | | ✅︎ |
409
| `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. | | ✅︎ |
410
| `GptOssForCausalLM` | GPT-OSS | `openai/gpt-oss-120b`, `openai/gpt-oss-20b` | ✅︎ | ✅︎ |
@@ -477,7 +476,6 @@ th {
476
| `SolarForCausalLM` | Solar Pro | `upstage/solar-pro-preview-instruct`, etc. | ✅︎ | ✅︎ |
477
| `StableLmForCausalLM` | StableLM | `stabilityai/stablelm-3b-4e1t`, `stabilityai/stablelm-base-alpha-7b-v2`, etc. | | |
478
| `StableLMEpochForCausalLM` | StableLM Epoch | `stabilityai/stablelm-zephyr-3b`, etc. | | ✅︎ |
480
-| `Starcoder2ForCausalLM` | Starcoder2 | `bigcode/starcoder2-3b`, `bigcode/starcoder2-7b`, `bigcode/starcoder2-15b`, etc. | | ✅︎ |
479
| `Step1ForCausalLM` | Step-Audio | `stepfun-ai/Step-Audio-EditX`, etc. | ✅︎ | ✅︎ |
480
| `Step3p5ForCausalLM` | Step-3.5-flash | `stepfun-ai/Step-3.5-Flash`, etc. | | ✅︎ |
481
| `TeleChatForCausalLM` | TeleChat | `chuhac/TeleChat2-35B`, etc. | ✅︎ | ✅︎ |
@@ -490,7 +488,9 @@ Some models are supported only via the [Transformers modeling backend](#transfor
488
489
| Architecture | Models | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) |
490
| ------------ | ------ | ----------------- | -------------------- | ------------------------- |
491
+| `GPTBigCodeForCausalLM` | StarCoder, SantaCoder, WizardCoder | `bigcode/starcoder`, `bigcode/gpt_bigcode-santacoder`, `WizardLM/WizardCoder-15B-V1.0`, etc. | ✅︎ | |
492
| `SmolLM3ForCausalLM` | SmolLM3 | `HuggingFaceTB/SmolLM3-3B` | ✅︎ | ✅︎ |
493
+| `Starcoder2ForCausalLM` | Starcoder2 | `bigcode/starcoder2-3b`, `bigcode/starcoder2-7b`, `bigcode/starcoder2-15b`, etc. | ✅︎ | ✅︎ |
494
495
!!! note
496
Currently, the ROCm version of vLLM supports Mistral and Mixtral only for context lengths up to 4096.
tests/distributed/test_pipeline_parallel.py
-1
@@ -115,7 +115,6 @@ TEXT_GENERATION_MODELS = {
115
"google/gemma-1.1-2b-it": PPTestSettings.fast(),
116
"google/gemma-2-9b": PPTestSettings.fast(),
117
"gpt2": PPTestSettings.fast(),
118
- "bigcode/starcoder": PPTestSettings.fast(),
118
"EleutherAI/gpt-j-6b": PPTestSettings.fast(),
119
"EleutherAI/pythia-1.4b": PPTestSettings.fast(),
120
"ibm/PowerLM-3b": PPTestSettings.fast(),
vllm/model_executor/models/gpt_bigcode.py
deleted
-339
@@ -1,339 +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/gpt2/modeling_gpt2.py
6
-# Copyright 2023 The vLLM team.
7
-# Copyright 2023 CTranslate2, and Michael Feil
8
-# Copyright 2018 The OpenAI Team Authors and HuggingFace Inc. team.
9
-# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
10
-#
11
-# Licensed under the Apache License, Version 2.0 (the "License");
12
-# you may not use this file except in compliance with the License.
13
-# You may obtain a copy of the License at
14
-#
15
-# http://www.apache.org/licenses/LICENSE-2.0
16
-#
17
-# Unless required by applicable law or agreed to in writing, software
18
-# distributed under the License is distributed on an "AS IS" BASIS,
19
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
20
-# See the License for the specific language governing permissions and
21
-# limitations under the License.
22
-"""Inference-only GPTBigCode model compatible with HuggingFace weights."""
23
-
24
-from collections.abc import Iterable
25
-from itertools import islice
26
-
27
-import torch
28
-from torch import nn
29
-from transformers import GPTBigCodeConfig
30
-
31
-from vllm.compilation.decorators import support_torch_compile
32
-from vllm.config import CacheConfig, VllmConfig
33
-from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
34
-from vllm.model_executor.layers.activation import get_act_fn
35
-from vllm.model_executor.layers.attention import Attention
36
-from vllm.model_executor.layers.linear import (
37
- ColumnParallelLinear,
38
- QKVParallelLinear,
39
- RowParallelLinear,
40
-)
41
-from vllm.model_executor.layers.logits_processor import LogitsProcessor
42
-from vllm.model_executor.layers.quantization import QuantizationConfig
43
-from vllm.model_executor.layers.vocab_parallel_embedding import (
44
- ParallelLMHead,
45
- VocabParallelEmbedding,
46
-)
47
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
48
-from vllm.sequence import IntermediateTensors
49
-
50
-from .interfaces import SupportsLoRA, SupportsPP
51
-from .utils import (
52
- AutoWeightsLoader,
53
- is_pp_missing_parameter,
54
- make_empty_intermediate_tensors_factory,
55
- make_layers,
56
- maybe_prefix,
57
-)
58
-
59
-
60
-class GPTBigCodeAttention(nn.Module):
61
- def __init__(
62
- self,
63
- config: GPTBigCodeConfig,
64
- cache_config: CacheConfig | None = None,
65
- quant_config: QuantizationConfig | None = None,
66
- prefix: str = "",
67
- ):
68
- super().__init__()
69
- self.hidden_size = config.hidden_size
70
- total_num_heads = config.num_attention_heads
71
- self.tensor_model_parallel_world_size = get_tensor_model_parallel_world_size()
72
- assert total_num_heads % self.tensor_model_parallel_world_size == 0
73
- self.num_heads = total_num_heads // self.tensor_model_parallel_world_size
74
- self.head_dim = self.hidden_size // total_num_heads
75
- self.scale = self.head_dim**-0.5
76
-
77
- self.multi_query = config.multi_query
78
- if self.multi_query:
79
- total_num_kv_heads = 1
80
- self.num_kv_heads = 1
81
- else:
82
- total_num_kv_heads = total_num_heads
83
- self.num_kv_heads = self.num_heads
84
- self.kv_dim = self.head_dim * self.num_kv_heads
85
- self.c_attn = QKVParallelLinear(
86
- self.hidden_size,
87
- self.head_dim,
88
- total_num_heads,
89
- total_num_kv_heads,
90
- bias=True,
91
- quant_config=quant_config,
92
- prefix=f"{prefix}.c_attn",
93
- )
94
-
95
- self.c_proj = RowParallelLinear(
96
- self.hidden_size,
97
- self.hidden_size,
98
- bias=True,
99
- quant_config=quant_config,
100
- prefix=f"{prefix}.c_proj",
101
- )
102
- self.attn = Attention(
103
- self.num_heads,
104
- self.head_dim,
105
- scale=self.scale,
106
- num_kv_heads=self.num_kv_heads,
107
- cache_config=cache_config,
108
- quant_config=quant_config,
109
- prefix=f"{prefix}.attn",
110
- )
111
-
112
- def forward(
113
- self,
114
- hidden_states: torch.Tensor,
115
- ) -> torch.Tensor:
116
- qkv, _ = self.c_attn(hidden_states)
117
- q, k, v = qkv.split(
118
- [
119
- self.hidden_size // self.tensor_model_parallel_world_size,
120
- self.kv_dim,
121
- self.kv_dim,
122
- ],
123
- dim=-1,
124
- )
125
- attn_output = self.attn(q, k, v)
126
- attn_output, _ = self.c_proj(attn_output)
127
- return attn_output
128
-
129
-
130
-class GPTBigMLP(nn.Module):
131
- def __init__(
132
- self,
133
- intermediate_size: int,
134
- config: GPTBigCodeConfig,
135
- quant_config: QuantizationConfig | None = None,
136
- prefix: str = "",
137
- ):
138
- super().__init__()
139
- hidden_size = config.hidden_size
140
- self.c_fc = ColumnParallelLinear(
141
- hidden_size,
142
- intermediate_size,
143
- bias=True,
144
- quant_config=quant_config,
145
- prefix=f"{prefix}.c_fc",
146
- )
147
- self.c_proj = RowParallelLinear(
148
- intermediate_size,
149
- hidden_size,
150
- bias=True,
151
- quant_config=quant_config,
152
- prefix=f"{prefix}.c_proj",
153
- )
154
- self.act = get_act_fn(config.activation_function)
155
-
156
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
157
- hidden_states, _ = self.c_fc(hidden_states)
158
- hidden_states = self.act(hidden_states)
159
- hidden_states, _ = self.c_proj(hidden_states)
160
- return hidden_states
161
-
162
-
163
-class GPTBigCodeBlock(nn.Module):
164
- def __init__(
165
- self,
166
- config: GPTBigCodeConfig,
167
- cache_config: CacheConfig | None = None,
168
- quant_config: QuantizationConfig | None = None,
169
- prefix: str = "",
170
- ):
171
- super().__init__()
172
- hidden_size = config.hidden_size
173
- inner_dim = config.n_inner if config.n_inner is not None else 4 * hidden_size
174
-
175
- self.ln_1 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
176
- self.attn = GPTBigCodeAttention(
177
- config, cache_config, quant_config, prefix=f"{prefix}.attn"
178
- )
179
- self.ln_2 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
180
- self.mlp = GPTBigMLP(inner_dim, config, quant_config, prefix=f"{prefix}.mlp")
181
-
182
- def forward(
183
- self,
184
- hidden_states: torch.Tensor,
185
- ) -> torch.Tensor:
186
- residual = hidden_states
187
- hidden_states = self.ln_1(hidden_states)
188
- attn_output = self.attn(
189
- hidden_states=hidden_states,
190
- )
191
- # residual connection
192
- hidden_states = attn_output + residual
193
-
194
- residual = hidden_states
195
- hidden_states = self.ln_2(hidden_states)
196
- feed_forward_hidden_states = self.mlp(hidden_states)
197
- # residual connection
198
- hidden_states = residual + feed_forward_hidden_states
199
- return hidden_states
200
-
201
-
202
-@support_torch_compile
203
-class GPTBigCodeModel(nn.Module):
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
- assert not config.add_cross_attention
213
-
214
- self.embed_dim = config.hidden_size
215
-
216
- self.vocab_size = config.vocab_size
217
- self.wte = VocabParallelEmbedding(
218
- self.vocab_size, self.embed_dim, org_num_embeddings=config.vocab_size
219
- )
220
- self.wpe = nn.Embedding(config.max_position_embeddings, self.embed_dim)
221
- self.start_layer, self.end_layer, self.h = make_layers(
222
- config.num_hidden_layers,
223
- lambda prefix: GPTBigCodeBlock(
224
- config, cache_config, quant_config, prefix=prefix
225
- ),
226
- prefix=f"{prefix}.h",
227
- )
228
- self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)
229
- self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
230
- ["hidden_states"], config.n_embd
231
- )
232
-
233
- def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
234
- return self.wte(input_ids)
235
-
236
- def forward(
237
- self,
238
- input_ids: torch.Tensor | None,
239
- position_ids: torch.Tensor,
240
- intermediate_tensors: IntermediateTensors | None,
241
- inputs_embeds: torch.Tensor | None = None,
242
- ) -> torch.Tensor | IntermediateTensors:
243
- if get_pp_group().is_first_rank:
244
- if inputs_embeds is None:
245
- inputs_embeds = self.embed_input_ids(input_ids)
246
- hidden_states = inputs_embeds + self.wpe(position_ids)
247
- else:
248
- hidden_states = intermediate_tensors["hidden_states"]
249
-
250
- for layer in islice(self.h, self.start_layer, self.end_layer):
251
- hidden_states = layer(hidden_states)
252
-
253
- if not get_pp_group().is_last_rank:
254
- return IntermediateTensors({"hidden_states": hidden_states})
255
- hidden_states = self.ln_f(hidden_states)
256
- return hidden_states
257
-
258
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
259
- params_dict = dict(self.named_parameters(remove_duplicate=False))
260
- loaded_params: set[str] = set()
261
- for name, loaded_weight in weights:
262
- if ".attn.bias" in name:
263
- # Skip attention mask.
264
- # NOTE: "c_attn.bias" should not be skipped.
265
- continue
266
- if is_pp_missing_parameter(name, self):
267
- continue
268
- param = params_dict[name]
269
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
270
- # TODO (@robertgshaw2-neuralmagic): move to fp8 linear method
271
- if "c_attn.input_scale" in name:
272
- weight_loader(param, loaded_weight, "q")
273
- weight_loader(param, loaded_weight, "k")
274
- weight_loader(param, loaded_weight, "v")
275
- else:
276
- weight_loader(param, loaded_weight)
277
- loaded_params.add(name)
278
- return loaded_params
279
-
280
-
281
-class GPTBigCodeForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
282
- packed_modules_mapping = {"c_attn": ["c_attn"]}
283
-
284
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
285
- super().__init__()
286
- config = vllm_config.model_config.hf_config
287
- quant_config = vllm_config.quant_config
288
-
289
- self.config = config
290
-
291
- self.quant_config = quant_config
292
- self.transformer = GPTBigCodeModel(
293
- vllm_config=vllm_config, prefix=maybe_prefix(prefix, "transformer")
294
- )
295
- if self.config.tie_word_embeddings:
296
- self.lm_head = self.transformer.wte
297
- else:
298
- self.lm_head = ParallelLMHead(
299
- self.transformer.vocab_size,
300
- self.transformer.embed_dim,
301
- prefix=maybe_prefix(prefix, "lm_head"),
302
- )
303
-
304
- self.logits_processor = LogitsProcessor(config.vocab_size)
305
- self.make_empty_intermediate_tensors = (
306
- self.transformer.make_empty_intermediate_tensors
307
- )
308
-
309
- def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
310
- return self.transformer.embed_input_ids(input_ids)
311
-
312
- def forward(
313
- self,
314
- input_ids: torch.Tensor | None,
315
- positions: torch.Tensor,
316
- intermediate_tensors: IntermediateTensors | None = None,
317
- inputs_embeds: torch.Tensor | None = None,
318
- ) -> torch.Tensor | IntermediateTensors:
319
- hidden_states = self.transformer(
320
- input_ids, positions, intermediate_tensors, inputs_embeds
321
- )
322
- return hidden_states
323
-
324
- def compute_logits(
325
- self,
326
- hidden_states: torch.Tensor,
327
- ) -> torch.Tensor | None:
328
- logits = self.logits_processor(self.lm_head, hidden_states)
329
- return logits
330
-
331
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
332
- skip_prefixes = None
333
- if self.config.tie_word_embeddings:
334
- skip_prefixes = ["lm_head."]
335
- loader = AutoWeightsLoader(
336
- self,
337
- skip_prefixes=skip_prefixes,
338
- )
339
- return loader.load_weights(weights)
vllm/model_executor/models/registry.py
+2
-2
@@ -116,7 +116,6 @@ _TEXT_GENERATION_MODELS = {
116
"GlmMoeDsaForCausalLM": ("deepseek_v2", "GlmMoeDsaForCausalLM"),
117
"GptOssForCausalLM": ("gpt_oss", "GptOssForCausalLM"),
118
"GPT2LMHeadModel": ("gpt2", "GPT2LMHeadModel"),
119
- "GPTBigCodeForCausalLM": ("gpt_bigcode", "GPTBigCodeForCausalLM"),
119
"GPTJForCausalLM": ("gpt_j", "GPTJForCausalLM"),
120
"GPTNeoXForCausalLM": ("gpt_neox", "GPTNeoXForCausalLM"),
121
"GraniteForCausalLM": ("granite", "GraniteForCausalLM"),
@@ -200,7 +199,6 @@ _TEXT_GENERATION_MODELS = {
199
"Step3p5ForCausalLM": ("step3p5", "Step3p5ForCausalLM"),
200
"StableLMEpochForCausalLM": ("stablelm", "StablelmForCausalLM"),
201
"StableLmForCausalLM": ("stablelm", "StablelmForCausalLM"),
203
- "Starcoder2ForCausalLM": ("starcoder2", "Starcoder2ForCausalLM"),
202
"SolarForCausalLM": ("solar", "SolarForCausalLM"),
203
"TeleChatForCausalLM": ("telechat2", "TeleChat2ForCausalLM"),
204
"TeleChat2ForCausalLM": ("telechat2", "TeleChat2ForCausalLM"),
@@ -630,7 +628,9 @@ _SPECULATIVE_DECODING_MODELS = {
628
629
_TRANSFORMERS_SUPPORTED_MODELS = {
630
# Text generation models
631
+ "GPTBigCodeForCausalLM": ("transformers", "TransformersForCausalLM"),
632
"SmolLM3ForCausalLM": ("transformers", "TransformersForCausalLM"),
633
+ "Starcoder2ForCausalLM": ("transformers", "TransformersForCausalLM"),
634
# Multimodal models
635
"Emu3ForConditionalGeneration": (
636
"transformers",
vllm/model_executor/models/starcoder2.py
deleted
-337
@@ -1,337 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-# Copyright 2024 BigCode and the HuggingFace Inc. team. All rights reserved.
5
-#
6
-# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
7
-# and OPT implementations in this library. It has been modified from its
8
-# original forms to accommodate minor architectural differences compared
9
-# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
10
-#
11
-# Licensed under the Apache License, Version 2.0 (the "License");
12
-# you may not use this file except in compliance with the License.
13
-# You may obtain a copy of the License at
14
-#
15
-# http://www.apache.org/licenses/LICENSE-2.0
16
-#
17
-# Unless required by applicable law or agreed to in writing, software
18
-# distributed under the License is distributed on an "AS IS" BASIS,
19
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
20
-# See the License for the specific language governing permissions and
21
-# limitations under the License.
22
-"""PyTorch Starcoder2 model."""
23
-
24
-from collections.abc import Iterable
25
-from itertools import islice
26
-
27
-import torch
28
-from torch import nn
29
-from transformers import Starcoder2Config
30
-
31
-from vllm.compilation.decorators import support_torch_compile
32
-from vllm.config import CacheConfig, VllmConfig
33
-from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
34
-from vllm.model_executor.layers.activation import get_act_fn
35
-from vllm.model_executor.layers.attention import Attention
36
-from vllm.model_executor.layers.linear import (
37
- ColumnParallelLinear,
38
- QKVParallelLinear,
39
- RowParallelLinear,
40
-)
41
-from vllm.model_executor.layers.logits_processor import LogitsProcessor
42
-from vllm.model_executor.layers.quantization import QuantizationConfig
43
-from vllm.model_executor.layers.rotary_embedding import get_rope
44
-from vllm.model_executor.layers.vocab_parallel_embedding import (
45
- ParallelLMHead,
46
- VocabParallelEmbedding,
47
-)
48
-from vllm.sequence import IntermediateTensors
49
-
50
-from .interfaces import SupportsPP
51
-from .utils import (
52
- AutoWeightsLoader,
53
- WeightsMapper,
54
- make_empty_intermediate_tensors_factory,
55
- make_layers,
56
- maybe_prefix,
57
-)
58
-
59
-
60
-class Starcoder2Attention(nn.Module):
61
- def __init__(
62
- self,
63
- config: Starcoder2Config,
64
- cache_config: CacheConfig | None = None,
65
- quant_config: QuantizationConfig | None = None,
66
- prefix: str = "",
67
- ):
68
- super().__init__()
69
- self.config = config
70
-
71
- self.hidden_size = config.hidden_size
72
- tp_size = get_tensor_model_parallel_world_size()
73
- self.total_num_heads = config.num_attention_heads
74
- assert self.total_num_heads % tp_size == 0
75
- self.num_heads = self.total_num_heads // tp_size
76
- self.total_num_kv_heads = config.num_key_value_heads
77
- if self.total_num_kv_heads >= tp_size:
78
- # Number of KV heads is greater than TP size, so we partition
79
- # the KV heads across multiple tensor parallel GPUs.
80
- assert self.total_num_kv_heads % tp_size == 0
81
- else:
82
- # Number of KV heads is less than TP size, so we replicate
83
- # the KV heads across multiple tensor parallel GPUs.
84
- assert tp_size % self.total_num_kv_heads == 0
85
- self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
86
- self.head_dim = self.hidden_size // self.total_num_heads
87
- self.q_size = self.num_heads * self.head_dim
88
- self.kv_size = self.num_kv_heads * self.head_dim
89
- self.scaling = self.head_dim**-0.5
90
- self.max_position_embeddings = config.max_position_embeddings
91
- self.use_bias = config.use_bias
92
-
93
- self.qkv_proj = QKVParallelLinear(
94
- self.hidden_size,
95
- self.head_dim,
96
- self.total_num_heads,
97
- self.total_num_kv_heads,
98
- bias=self.use_bias,
99
- quant_config=quant_config,
100
- prefix=f"{prefix}.qkv_proj",
101
- )
102
- self.o_proj = RowParallelLinear(
103
- self.total_num_heads * self.head_dim,
104
- self.hidden_size,
105
- bias=self.use_bias,
106
- quant_config=quant_config,
107
- prefix=f"{prefix}.o_proj",
108
- )
109
- self.rotary_emb = get_rope(
110
- self.head_dim,
111
- max_position=self.max_position_embeddings,
112
- rope_parameters=config.rope_parameters,
113
- is_neox_style=True,
114
- )
115
- self.attn = Attention(
116
- self.num_heads,
117
- self.head_dim,
118
- self.scaling,
119
- num_kv_heads=self.num_kv_heads,
120
- cache_config=cache_config,
121
- quant_config=quant_config,
122
- prefix=f"{prefix}.attn",
123
- )
124
-
125
- def forward(
126
- self,
127
- positions: torch.Tensor,
128
- hidden_states: torch.Tensor,
129
- ) -> torch.Tensor:
130
- qkv, _ = self.qkv_proj(hidden_states)
131
- q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
132
- q, k = self.rotary_emb(positions, q, k)
133
- attn_output = self.attn(q, k, v)
134
- output, _ = self.o_proj(attn_output)
135
- return output
136
-
137
-
138
-class Starcoder2MLP(nn.Module):
139
- def __init__(
140
- self,
141
- config: Starcoder2Config,
142
- quant_config: QuantizationConfig | None = None,
143
- prefix: str = "",
144
- ):
145
- super().__init__()
146
- self.c_fc = ColumnParallelLinear(
147
- config.hidden_size,
148
- config.intermediate_size,
149
- bias=config.use_bias,
150
- quant_config=quant_config,
151
- prefix=f"{prefix}.c_fc",
152
- )
153
- self.c_proj = RowParallelLinear(
154
- config.intermediate_size,
155
- config.hidden_size,
156
- bias=config.use_bias,
157
- quant_config=quant_config,
158
- prefix=f"{prefix}.c_proj",
159
- )
160
- self.act = get_act_fn(config.hidden_act)
161
-
162
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
163
- hidden_states, _ = self.c_fc(hidden_states)
164
- hidden_states = self.act(hidden_states)
165
- hidden_states, _ = self.c_proj(hidden_states)
166
- return hidden_states
167
-
168
-
169
-class Starcoder2DecoderLayer(nn.Module):
170
- def __init__(
171
- self,
172
- config: Starcoder2Config,
173
- cache_config: CacheConfig | None = None,
174
- quant_config: QuantizationConfig | None = None,
175
- prefix: str = "",
176
- ):
177
- super().__init__()
178
- self.hidden_size = config.hidden_size
179
- self.self_attn = Starcoder2Attention(
180
- config,
181
- cache_config,
182
- quant_config=quant_config,
183
- prefix=f"{prefix}.self_attn",
184
- )
185
- self.mlp = Starcoder2MLP(
186
- config, quant_config=quant_config, prefix=f"{prefix}.mlp"
187
- )
188
- self.input_layernorm = nn.LayerNorm(config.hidden_size, eps=config.norm_epsilon)
189
- self.post_attention_layernorm = nn.LayerNorm(
190
- config.hidden_size, eps=config.norm_epsilon
191
- )
192
-
193
- def forward(
194
- self,
195
- positions: torch.Tensor,
196
- hidden_states: torch.Tensor,
197
- ) -> torch.Tensor:
198
- # Self Attention
199
- residual = hidden_states
200
- hidden_states = self.input_layernorm(hidden_states)
201
- hidden_states = self.self_attn(
202
- positions=positions,
203
- hidden_states=hidden_states,
204
- )
205
- hidden_states = residual + hidden_states
206
-
207
- # Fully Connected
208
- residual = hidden_states
209
- hidden_states = self.post_attention_layernorm(hidden_states)
210
- hidden_states = self.mlp(hidden_states)
211
- hidden_states = residual + hidden_states
212
-
213
- return hidden_states
214
-
215
-
216
-@support_torch_compile
217
-class Starcoder2Model(nn.Module):
218
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
219
- super().__init__()
220
-
221
- config = vllm_config.model_config.hf_config
222
- cache_config = vllm_config.cache_config
223
- quant_config = vllm_config.quant_config
224
-
225
- self.config = config
226
- self.vocab_size = config.vocab_size
227
-
228
- self.embed_tokens = VocabParallelEmbedding(
229
- config.vocab_size,
230
- config.hidden_size,
231
- quant_config=quant_config,
232
- prefix=f"{prefix}.embed_tokens",
233
- )
234
- self.start_layer, self.end_layer, self.layers = make_layers(
235
- config.num_hidden_layers,
236
- lambda prefix: Starcoder2DecoderLayer(
237
- config, cache_config, quant_config=quant_config, prefix=prefix
238
- ),
239
- prefix=f"{prefix}.layers",
240
- )
241
- self.norm = nn.LayerNorm(config.hidden_size, eps=config.norm_epsilon)
242
- self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
243
- ["hidden_states"], config.hidden_size
244
- )
245
-
246
- def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
247
- return self.embed_tokens(input_ids)
248
-
249
- def forward(
250
- self,
251
- input_ids: torch.Tensor | None,
252
- positions: torch.Tensor,
253
- intermediate_tensors: IntermediateTensors | None,
254
- inputs_embeds: torch.Tensor | None = None,
255
- ) -> torch.Tensor | IntermediateTensors:
256
- if get_pp_group().is_first_rank:
257
- if inputs_embeds is not None:
258
- hidden_states = inputs_embeds
259
- else:
260
- hidden_states = self.embed_input_ids(input_ids)
261
- else:
262
- assert intermediate_tensors is not None
263
- hidden_states = intermediate_tensors["hidden_states"]
264
- for layer in islice(self.layers, self.start_layer, self.end_layer):
265
- hidden_states = layer(positions, hidden_states)
266
- if not get_pp_group().is_last_rank:
267
- return IntermediateTensors({"hidden_states": hidden_states})
268
- hidden_states = self.norm(hidden_states)
269
- return hidden_states
270
-
271
-
272
-class Starcoder2ForCausalLM(nn.Module, SupportsPP):
273
- hf_to_vllm_mapper = WeightsMapper(
274
- orig_to_new_stacked={
275
- # weight_name: (param_name, shard_id)
276
- ".q_proj": (".qkv_proj", "q"),
277
- ".k_proj": (".qkv_proj", "k"),
278
- ".v_proj": (".qkv_proj", "v"),
279
- }
280
- )
281
-
282
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
283
- super().__init__()
284
- config = vllm_config.model_config.hf_config
285
- quant_config = vllm_config.quant_config
286
- self.config = config
287
- self.model = Starcoder2Model(
288
- vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
289
- )
290
- self.vocab_size = config.vocab_size
291
-
292
- if config.tie_word_embeddings:
293
- self.lm_head = self.model.embed_tokens
294
- else:
295
- self.lm_head = ParallelLMHead(
296
- config.vocab_size,
297
- config.hidden_size,
298
- quant_config=quant_config,
299
- prefix=f"{prefix}.lm_head",
300
- )
301
- self.logits_processor = LogitsProcessor(config.vocab_size)
302
- self.make_empty_intermediate_tensors = (
303
- self.model.make_empty_intermediate_tensors
304
- )
305
-
306
- def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
307
- return self.model.embed_input_ids(input_ids)
308
-
309
- def forward(
310
- self,
311
- input_ids: torch.Tensor | None,
312
- positions: torch.Tensor,
313
- intermediate_tensors: IntermediateTensors | None = None,
314
- inputs_embeds: torch.Tensor | None = None,
315
- ) -> torch.Tensor | IntermediateTensors:
316
- hidden_states = self.model(
317
- input_ids, positions, intermediate_tensors, inputs_embeds
318
- )
319
- return hidden_states
320
-
321
- def compute_logits(
322
- self,
323
- hidden_states: torch.Tensor,
324
- ) -> torch.Tensor | None:
325
- logits = self.logits_processor(self.lm_head, hidden_states)
326
- return logits
327
-
328
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
329
- loader = AutoWeightsLoader(
330
- self,
331
- # Models trained using ColossalAI may include these tensors in
332
- # the checkpoint. Skip them.
333
- skip_prefixes=(
334
- ["lm_head.weight"] if self.config.tie_word_embeddings else None
335
- ),
336
- )
337
- return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)