gpt_neox.py 13 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0

3
4
# Adapted from
# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/gpt_neox/modeling_gpt_neox.py
Woosuk Kwon's avatar
Woosuk Kwon committed
5
# Copyright 2023 The vLLM team.
6
7
8
9
10
11
12
13
14
15
16
17
18
# Copyright 2022 EleutherAI The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
Woosuk Kwon's avatar
Woosuk Kwon committed
19
"""Inference-only GPT-NeoX model compatible with HuggingFace weights."""
20
from typing import Iterable, Optional, Set, Tuple, Union
21
22
23

import torch
from torch import nn
24
25
from transformers import GPTNeoXConfig

26
from vllm.attention import Attention
27
from vllm.compilation.decorators import support_torch_compile
28
from vllm.config import CacheConfig, VllmConfig
29
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
Woosuk Kwon's avatar
Woosuk Kwon committed
30
from vllm.model_executor.layers.activation import get_act_fn
31
32
33
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
                                               QKVParallelLinear,
                                               RowParallelLinear)
34
from vllm.model_executor.layers.logits_processor import LogitsProcessor
35
from vllm.model_executor.layers.quantization import QuantizationConfig
36
from vllm.model_executor.layers.rotary_embedding import get_rope
37
from vllm.model_executor.layers.vocab_parallel_embedding import (
38
    ParallelLMHead, VocabParallelEmbedding)
39
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
40
from vllm.model_executor.sampling_metadata import SamplingMetadata
41
from vllm.sequence import IntermediateTensors
42

43
from .interfaces import SupportsPP
44
from .utils import (AutoWeightsLoader, is_pp_missing_parameter,
45
46
                    make_empty_intermediate_tensors_factory, make_layers,
                    maybe_prefix)
47

48
49
50

class GPTNeoXAttention(nn.Module):

51
52
53
    def __init__(
        self,
        config: GPTNeoXConfig,
54
        cache_config: Optional[CacheConfig] = None,
55
        quant_config: Optional[QuantizationConfig] = None,
56
        prefix: str = "",
57
    ):
58
59
60
61
        super().__init__()
        self.total_num_heads = config.num_attention_heads
        self.hidden_size = config.hidden_size
        self.head_size = self.hidden_size // self.total_num_heads
62
        self.bias = getattr(config, "attention_bias", True)
63

64
65
        tensor_model_parallel_world_size = (
            get_tensor_model_parallel_world_size())
66
        assert self.total_num_heads % tensor_model_parallel_world_size == 0
67
68
69
        self.num_heads = (self.total_num_heads //
                          tensor_model_parallel_world_size)

70
        self.query_key_value = QKVParallelLinear(
71
            config.hidden_size,
72
73
            self.head_size,
            self.total_num_heads,
74
            bias=self.bias,
75
            quant_config=quant_config,
76
77
78
79
        )
        self.dense = RowParallelLinear(
            config.hidden_size,
            config.hidden_size,
80
            bias=self.bias,
81
            quant_config=quant_config,
82
        )
83
        scaling = self.head_size**-0.5
84
85
        rotary_dim = int(self.head_size * config.rotary_pct)
        assert rotary_dim % 2 == 0
86
87
88
        rope_theta = getattr(config, "rope_theta", 10000)
        max_position_embeddings = getattr(config, "max_position_embeddings",
                                          8192)
Woosuk Kwon's avatar
Woosuk Kwon committed
89
        self.rotary_emb = get_rope(
90
            self.head_size,
Woosuk Kwon's avatar
Woosuk Kwon committed
91
92
            rotary_dim=rotary_dim,
            max_position=max_position_embeddings,
93
            base=rope_theta,
Woosuk Kwon's avatar
Woosuk Kwon committed
94
        )
95
96
97
        self.attn = Attention(self.num_heads,
                              self.head_size,
                              scaling,
98
                              cache_config=cache_config,
99
100
                              quant_config=quant_config,
                              prefix=f"{prefix}.attn")
101
102
103

    def forward(
        self,
104
        position_ids: torch.Tensor,
105
106
107
108
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        qkv, _ = self.query_key_value(hidden_states)
        q, k, v = qkv.chunk(chunks=3, dim=-1)
Woosuk Kwon's avatar
Woosuk Kwon committed
109
        q, k = self.rotary_emb(position_ids, q, k)
110
        attn_output = self.attn(q, k, v)
111
112
113
114
115
        output, _ = self.dense(attn_output)
        return output


class GPTNeoXMLP(nn.Module):
116

117
118
119
    def __init__(
        self,
        config: GPTNeoXConfig,
120
        quant_config: Optional[QuantizationConfig] = None,
121
    ):
122
        super().__init__()
123
124
125
        self.dense_h_to_4h = ColumnParallelLinear(
            config.hidden_size,
            config.intermediate_size,
126
            quant_config=quant_config,
127
128
129
130
        )
        self.dense_4h_to_h = RowParallelLinear(
            config.intermediate_size,
            config.hidden_size,
131
            quant_config=quant_config,
132
        )
133
        self.act = get_act_fn(config.hidden_act)
134
135
136
137
138
139
140
141
142
143

    def forward(self, hidden_states):
        hidden_states, _ = self.dense_h_to_4h(hidden_states)
        hidden_states = self.act(hidden_states)
        hidden_states, _ = self.dense_4h_to_h(hidden_states)
        return hidden_states


class GPTNeoXLayer(nn.Module):

144
145
146
    def __init__(
        self,
        config: GPTNeoXConfig,
147
        cache_config: Optional[CacheConfig] = None,
148
        quant_config: Optional[QuantizationConfig] = None,
149
        prefix: str = "",
150
    ):
151
152
        super().__init__()
        self.use_parallel_residual = config.use_parallel_residual
153
154
155
156
        self.input_layernorm = nn.LayerNorm(config.hidden_size,
                                            eps=config.layer_norm_eps)
        self.post_attention_layernorm = nn.LayerNorm(config.hidden_size,
                                                     eps=config.layer_norm_eps)
157
158
159
160
        self.attention = GPTNeoXAttention(config,
                                          cache_config,
                                          quant_config,
                                          prefix=f"{prefix}.attention")
161
        self.mlp = GPTNeoXMLP(config, quant_config)
162
163
164

    def forward(
        self,
165
        position_ids: torch.Tensor,
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        attn_input = self.input_layernorm(hidden_states)
        attn_output = self.attention(
            position_ids=position_ids,
            hidden_states=attn_input,
        )

        if self.use_parallel_residual:
            # pseudocode:
            # x = x + attn(ln1(x)) + mlp(ln2(x))
            mlp_input = self.post_attention_layernorm(hidden_states)
            mlp_output = self.mlp(mlp_input)
            hidden_states = mlp_output + attn_output + hidden_states
        else:
            # pseudocode:
            # x = x + attn(ln1(x))
            # x = x + mlp(ln2(x))
            attn_output = attn_output + hidden_states
            mlp_input = self.post_attention_layernorm(attn_output)
            mlp_output = self.mlp(mlp_input)
            hidden_states = mlp_output + attn_output
        return hidden_states


191
@support_torch_compile
192
class GPTNeoXModel(nn.Module):
Woosuk Kwon's avatar
Woosuk Kwon committed
193

194
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
195
        super().__init__()
196
197
198
199
200

        config = vllm_config.model_config.hf_config
        cache_config = vllm_config.cache_config
        quant_config = vllm_config.quant_config

201
202
        self.config = config

203
204
205
206
        self.embed_in = VocabParallelEmbedding(
            config.vocab_size,
            config.hidden_size,
        )
207
208
        self.start_layer, self.end_layer, self.layers = make_layers(
            config.num_hidden_layers,
209
210
            lambda prefix: GPTNeoXLayer(
                config, cache_config, quant_config, prefix=prefix),
211
212
            prefix=f"{prefix}.layers",
        )
213
214
        self.final_layer_norm = nn.LayerNorm(config.hidden_size,
                                             eps=config.layer_norm_eps)
215
216
217
        self.make_empty_intermediate_tensors = (
            make_empty_intermediate_tensors_factory(["hidden_states"],
                                                    config.hidden_size))
218

219
220
221
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embed_in(input_ids)

222
223
    def forward(
        self,
224
225
        input_ids: torch.Tensor,
        position_ids: torch.Tensor,
226
        intermediate_tensors: Optional[IntermediateTensors],
227
        inputs_embeds: Optional[torch.Tensor] = None,
228
229
    ) -> Union[torch.Tensor, IntermediateTensors]:
        if get_pp_group().is_first_rank:
230
231
232
233
            if inputs_embeds is not None:
                hidden_states = inputs_embeds
            else:
                hidden_states = self.get_input_embeddings(input_ids)
234
235
        else:
            hidden_states = intermediate_tensors["hidden_states"]
236
237
        for layer in self.layers[self.start_layer:self.end_layer]:
            hidden_states = layer(position_ids, hidden_states)
238
239
        if not get_pp_group().is_last_rank:
            return IntermediateTensors({"hidden_states": hidden_states})
240
241
242
        hidden_states = self.final_layer_norm(hidden_states)
        return hidden_states

243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
        params_dict = dict(self.named_parameters())
        loaded_params: Set[str] = set()
        for name, loaded_weight in weights:
            if ("attention.bias" in name or "attention.masked_bias" in name
                    or "rotary_emb.inv_freq" in name):
                continue
            if ("rotary_emb.cos_cached" in name
                    or "rotary_emb.sin_cached" in name):
                # Models trained using OpenRLHF may include
                # these tensors in the checkpoint. Skip them.
                continue
            if is_pp_missing_parameter(name, self):
                continue
            param = params_dict[name]

            if "query_key_value" in name:
                # NOTE: GPT-NeoX's fused QKV's output_dim has the shape of
                # (num_heads * 3 * head_size), while the
                # required shape is (3 * num_heads * head_size).
                # Thus, we need weight conversion.
                output_dim = getattr(param, "output_dim", None)
                num_heads = self.config.num_attention_heads
                if output_dim is not None:
                    loaded_weight_shape = loaded_weight.shape
                    loaded_weight = loaded_weight.view(
                        loaded_weight_shape[:output_dim] + (num_heads, 3, -1) +
                        loaded_weight_shape[output_dim + 1:])
                    loaded_weight = loaded_weight.transpose(
                        output_dim, output_dim + 1)
                    loaded_weight = loaded_weight.reshape(loaded_weight_shape)

            weight_loader = getattr(param, "weight_loader",
                                    default_weight_loader)
            weight_loader(param, loaded_weight)
            loaded_params.add(name)
        return loaded_params

282

283
class GPTNeoXForCausalLM(nn.Module, SupportsPP):
284

285
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
286
        super().__init__()
287
288
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
289
        self.config = config
290
        self.quant_config = quant_config
291
292
        self.gpt_neox = GPTNeoXModel(vllm_config=vllm_config,
                                     prefix=maybe_prefix(prefix, "gpt_neox"))
293
        self.embed_out = ParallelLMHead(
294
            config.vocab_size,
295
            config.hidden_size,
296
            quant_config=quant_config,
297
        )
298
299
        if self.config.tie_word_embeddings:
            self.embed_out.weight = self.gpt_neox.embed_in.weight
300
        self.logits_processor = LogitsProcessor(config.vocab_size)
301
302
        self.make_empty_intermediate_tensors = (
            self.gpt_neox.make_empty_intermediate_tensors)
303

304
305
306
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.gpt_neox.get_input_embeddings(input_ids)

307
308
    def forward(
        self,
309
310
        input_ids: torch.Tensor,
        positions: torch.Tensor,
311
        intermediate_tensors: Optional[IntermediateTensors] = None,
312
        inputs_embeds: Optional[torch.Tensor] = None,
313
    ) -> Union[torch.Tensor, IntermediateTensors]:
314
315
        hidden_states = self.gpt_neox(input_ids, positions,
                                      intermediate_tensors, inputs_embeds)
316
317
        return hidden_states

318
319
320
321
322
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
323
        logits = self.logits_processor(self.embed_out, hidden_states,
324
325
326
                                       sampling_metadata)
        return logits

327
328
    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
329
330
        loader = AutoWeightsLoader(self)
        return loader.load_weights(weights)