gpt_j.py 11.6 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
# coding=utf-8
# Adapted from
# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/gptj/modeling_gptj.py
# Copyright 2023 The vLLM team.
# Copyright 2021 The EleutherAI and HuggingFace Teams. 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
18
"""Inference-only GPT-J model compatible with HuggingFace weights."""
19
from typing import Iterable, List, Optional, Tuple, Union
20
21
22
23
24

import torch
from torch import nn
from transformers import GPTJConfig

25
from vllm.attention import Attention, AttentionMetadata
26
from vllm.config import CacheConfig
27
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
28
from vllm.model_executor.layers.activation import get_act_fn
29
30
31
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
                                               QKVParallelLinear,
                                               RowParallelLinear)
32
from vllm.model_executor.layers.logits_processor import LogitsProcessor
33
from vllm.model_executor.layers.quantization import QuantizationConfig
34
from vllm.model_executor.layers.rotary_embedding import get_rope
35
from vllm.model_executor.layers.sampler import Sampler, SamplerOutput
36
from vllm.model_executor.layers.vocab_parallel_embedding import (
37
    ParallelLMHead, VocabParallelEmbedding)
38
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
39
from vllm.model_executor.sampling_metadata import SamplingMetadata
40
from vllm.sequence import IntermediateTensors
41

42
43
44
45
from .interfaces import SupportsPP
from .utils import (is_pp_missing_parameter,
                    make_empty_intermediate_tensors_factory, make_layers)

46
47
48

class GPTJAttention(nn.Module):

49
50
51
    def __init__(
        self,
        config: GPTJConfig,
52
        cache_config: Optional[CacheConfig] = None,
53
        quant_config: Optional[QuantizationConfig] = None,
54
    ):
55
56
57
58
59
        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

60
        self.qkv_proj = QKVParallelLinear(
61
            config.hidden_size,
62
63
            self.head_size,
            self.total_num_heads,
64
            bias=False,
65
            quant_config=quant_config,
66
67
68
69
70
        )
        self.out_proj = RowParallelLinear(
            config.hidden_size,
            config.hidden_size,
            bias=False,
71
            quant_config=quant_config,
72
        )
73
74
75
76
77
78

        tp_world_size = get_tensor_model_parallel_world_size()
        assert self.total_num_heads % tp_world_size == 0
        self.num_heads = self.total_num_heads // tp_world_size

        scaling = self.head_size**-0.5
79
        assert getattr(config, "rotary", True)
80
        assert config.rotary_dim % 2 == 0
81
82
83
        rope_theta = getattr(config, "rope_theta", 10000)
        max_position_embeddings = getattr(config, "max_position_embeddings",
                                          8192)
Woosuk Kwon's avatar
Woosuk Kwon committed
84
        self.rotary_emb = get_rope(
85
            self.head_size,
Woosuk Kwon's avatar
Woosuk Kwon committed
86
            rotary_dim=config.rotary_dim,
87
            max_position=max_position_embeddings,
Woosuk Kwon's avatar
Woosuk Kwon committed
88
89
90
            base=rope_theta,
            is_neox_style=False,
        )
91
92
93
        self.attn = Attention(self.num_heads,
                              self.head_size,
                              scaling,
94
95
                              cache_config=cache_config,
                              quant_config=quant_config)
96
97
98
99
100

    def forward(
        self,
        position_ids: torch.Tensor,
        hidden_states: torch.Tensor,
101
102
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
103
104
105
    ) -> torch.Tensor:
        qkv, _ = self.qkv_proj(hidden_states)
        q, k, v = qkv.chunk(chunks=3, dim=-1)
Woosuk Kwon's avatar
Woosuk Kwon committed
106
        q, k = self.rotary_emb(position_ids, q, k)
107
        attn_output = self.attn(q, k, v, kv_cache, attn_metadata)
108
109
110
111
112
113
        attn_output, _ = self.out_proj(attn_output)
        return attn_output


class GPTJMLP(nn.Module):

114
115
116
117
    def __init__(
        self,
        intermediate_size: int,
        config: GPTJConfig,
118
        quant_config: Optional[QuantizationConfig] = None,
119
    ):
120
121
        super().__init__()
        hidden_size = config.n_embd
122
123
124
        self.fc_in = ColumnParallelLinear(
            hidden_size,
            intermediate_size,
125
            quant_config=quant_config,
126
127
128
129
        )
        self.fc_out = RowParallelLinear(
            intermediate_size,
            hidden_size,
130
            quant_config=quant_config,
131
        )
132
133
        self.act = get_act_fn(config.activation_function, quant_config,
                              intermediate_size)
134
135
136
137
138
139
140
141
142
143

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        hidden_states, _ = self.fc_in(hidden_states)
        hidden_states = self.act(hidden_states)
        hidden_states, _ = self.fc_out(hidden_states)
        return hidden_states


class GPTJBlock(nn.Module):

144
145
146
    def __init__(
        self,
        config: GPTJConfig,
147
        cache_config: Optional[CacheConfig] = None,
148
        quant_config: Optional[QuantizationConfig] = None,
149
    ):
150
        super().__init__()
151
152
        inner_dim = (4 * config.n_embd
                     if config.n_inner is None else config.n_inner)
153
        self.ln_1 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
154
        self.attn = GPTJAttention(config, cache_config, quant_config)
155
        self.mlp = GPTJMLP(inner_dim, config, quant_config)
156
157
158
159
160

    def forward(
        self,
        position_ids: torch.Tensor,
        hidden_states: torch.Tensor,
161
162
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
163
164
165
166
167
168
169
    ) -> torch.Tensor:
        residual = hidden_states
        hidden_states = self.ln_1(hidden_states)
        attn_output = self.attn(
            position_ids=position_ids,
            hidden_states=hidden_states,
            kv_cache=kv_cache,
170
            attn_metadata=attn_metadata,
171
172
173
174
175
176
177
178
        )
        mlp_output = self.mlp(hidden_states)
        hidden_states = attn_output + mlp_output + residual
        return hidden_states


class GPTJModel(nn.Module):

179
180
181
    def __init__(
        self,
        config: GPTJConfig,
182
        cache_config: Optional[CacheConfig] = None,
183
        quant_config: Optional[QuantizationConfig] = None,
184
        prefix: str = "",
185
    ):
186
187
188
        super().__init__()
        self.config = config
        self.embed_dim = config.n_embd
189
190
191
192
        self.wte = VocabParallelEmbedding(
            config.vocab_size,
            self.embed_dim,
        )
193
194
195
196
197
        self.start_layer, self.end_layer, self.h = make_layers(
            config.n_layer,
            lambda prefix: GPTJBlock(config, cache_config, quant_config),
            prefix=f"{prefix}.h",
        )
198
        self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)
199
200
201
        self.make_empty_intermediate_tensors = (
            make_empty_intermediate_tensors_factory(["hidden_states"],
                                                    config.n_embd))
202
203
204
205
206

    def forward(
        self,
        input_ids: torch.Tensor,
        position_ids: torch.Tensor,
207
208
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
209
210
211
212
213
214
215
        intermediate_tensors: Optional[IntermediateTensors],
    ) -> Union[torch.Tensor, IntermediateTensors]:
        if get_pp_group().is_first_rank:
            hidden_states = self.wte(input_ids)
        else:
            hidden_states = intermediate_tensors["hidden_states"]
        for i in range(self.start_layer, self.end_layer):
216
217
218
219
            layer = self.h[i]
            hidden_states = layer(
                position_ids,
                hidden_states,
220
                kv_caches[i - self.start_layer],
221
                attn_metadata,
222
            )
223
224
        if not get_pp_group().is_last_rank:
            return IntermediateTensors({"hidden_states": hidden_states})
225
226
227
228
        hidden_states = self.ln_f(hidden_states)
        return hidden_states


229
class GPTJForCausalLM(nn.Module, SupportsPP):
230

231
232
233
    def __init__(
        self,
        config: GPTJConfig,
234
        cache_config: Optional[CacheConfig] = None,
235
        quant_config: Optional[QuantizationConfig] = None,
236
    ):
237
238
        super().__init__()
        self.config = config
239
        self.quant_config = quant_config
240
        assert not config.tie_word_embeddings
241
        self.transformer = GPTJModel(config, cache_config, quant_config)
242
        self.lm_head = ParallelLMHead(
243
            config.vocab_size,
244
245
            config.n_embd,
            bias=True,
246
            quant_config=quant_config,
247
        )
248
249
        self.logits_processor = LogitsProcessor(config.vocab_size)
        self.sampler = Sampler()
250
251
        self.make_empty_intermediate_tensors = (
            self.transformer.make_empty_intermediate_tensors)
252
253
254
255
256

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
257
258
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
259
        intermediate_tensors: Optional[IntermediateTensors] = None,
260
    ) -> Union[torch.Tensor, IntermediateTensors]:
261
        hidden_states = self.transformer(input_ids, positions, kv_caches,
262
                                         attn_metadata, intermediate_tensors)
263
264
        return hidden_states

265
266
267
268
269
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
270
        logits = self.logits_processor(self.lm_head, hidden_states,
271
272
273
                                       sampling_metadata, self.lm_head.bias)
        return logits

274
275
    def sample(
        self,
276
        logits: torch.Tensor,
277
        sampling_metadata: SamplingMetadata,
278
    ) -> Optional[SamplerOutput]:
279
        next_tokens = self.sampler(logits, sampling_metadata)
280
281
        return next_tokens

282
    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
283
284
285
286
287
288
289
290
291
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
            ("qkv_proj", "q_proj", "q"),
            ("qkv_proj", "k_proj", "k"),
            ("qkv_proj", "v_proj", "v"),
            ("gate_up_proj", "gate_proj", 0),
            ("gate_up_proj", "up_proj", 1),
        ]
        params_dict = dict(self.named_parameters())
292
        for name, loaded_weight in weights:
293
294
            if "attn.bias" in name or "attn.masked_bias" in name:
                continue
295
296
            for (param_name, weight_name, shard_id) in stacked_params_mapping:
                if weight_name not in name:
297
                    continue
CHU Tianxiang's avatar
CHU Tianxiang committed
298
299
300
301
                name = name.replace(weight_name, param_name)
                # Skip loading extra bias for GPTQ models.
                if name.endswith(".bias") and name not in params_dict:
                    continue
302
303
                if is_pp_missing_parameter(name, self):
                    continue
CHU Tianxiang's avatar
CHU Tianxiang committed
304
                param = params_dict[name]
305
306
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
307
                break
308
            else:
CHU Tianxiang's avatar
CHU Tianxiang committed
309
310
311
                # Skip loading extra bias for GPTQ models.
                if name.endswith(".bias") and name not in params_dict:
                    continue
312
313
                if is_pp_missing_parameter(name, self):
                    continue
314
315
316
317
                param = params_dict[name]
                weight_loader = getattr(param, "weight_loader",
                                        default_weight_loader)
                weight_loader(param, loaded_weight)