gpt_neox.py 15.7 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

gaoqiong's avatar
gaoqiong committed
22
23
import os
import re
24
25
import torch
from torch import nn
26
27
from transformers import GPTNeoXConfig

28
from vllm.attention import Attention
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
Woosuk Kwon's avatar
Woosuk Kwon committed
32
from vllm.model_executor.layers.activation import get_act_fn
33
34
35
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
                                               QKVParallelLinear,
                                               RowParallelLinear)
36
from vllm.model_executor.layers.logits_processor import LogitsProcessor
37
from vllm.model_executor.layers.quantization import QuantizationConfig
38
from vllm.model_executor.layers.rotary_embedding import get_rope
Joe Runde's avatar
Joe Runde committed
39
from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
40
from vllm.model_executor.layers.vocab_parallel_embedding import (
41
    ParallelLMHead, VocabParallelEmbedding)
42
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
43
from vllm.model_executor.sampling_metadata import SamplingMetadata
44
from vllm.sequence import IntermediateTensors
45

46
47
from .interfaces import SupportsPP
from .utils import (is_pp_missing_parameter,
48
49
                    make_empty_intermediate_tensors_factory, make_layers,
                    maybe_prefix)
gaoqiong's avatar
gaoqiong committed
50
51
from vllm.utils import is_hip,W8a8GetCacheJSON
from vllm import _custom_ops as ops
52
53
54

class GPTNeoXAttention(nn.Module):

55
56
57
    def __init__(
        self,
        config: GPTNeoXConfig,
58
        cache_config: Optional[CacheConfig] = None,
59
        quant_config: Optional[QuantizationConfig] = None,
60
        prefix: str = "",
61
    ):
62
63
64
65
        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
66
        self.bias = getattr(config, "attention_bias", True)
67

68
69
        tensor_model_parallel_world_size = (
            get_tensor_model_parallel_world_size())
70
        assert self.total_num_heads % tensor_model_parallel_world_size == 0
71
72
73
        self.num_heads = (self.total_num_heads //
                          tensor_model_parallel_world_size)

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

    def forward(
        self,
108
        position_ids: torch.Tensor,
109
110
111
112
        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
113
        q, k = self.rotary_emb(position_ids, q, k)
114
        attn_output = self.attn(q, k, v)
115
116
117
118
119
        output, _ = self.dense(attn_output)
        return output


class GPTNeoXMLP(nn.Module):
120

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

    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):

148
149
150
    def __init__(
        self,
        config: GPTNeoXConfig,
151
        cache_config: Optional[CacheConfig] = None,
152
        quant_config: Optional[QuantizationConfig] = None,
153
        prefix: str = "",
154
    ):
155
156
        super().__init__()
        self.use_parallel_residual = config.use_parallel_residual
157
158
159
160
        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)
161
162
163
164
        self.attention = GPTNeoXAttention(config,
                                          cache_config,
                                          quant_config,
                                          prefix=f"{prefix}.attention")
165
        self.mlp = GPTNeoXMLP(config, quant_config)
166
167
168

    def forward(
        self,
169
        position_ids: torch.Tensor,
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
        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


195
@support_torch_compile
196
class GPTNeoXModel(nn.Module):
Woosuk Kwon's avatar
Woosuk Kwon committed
197

198
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
199
        super().__init__()
200
201
202
203
204

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

205
206
        self.config = config

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

223
224
225
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embed_in(input_ids)

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


248
class GPTNeoXForCausalLM(nn.Module, SupportsPP):
249

250
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
251
        super().__init__()
252
253
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
254
        self.config = config
255
        self.quant_config = quant_config
256
257
        self.gpt_neox = GPTNeoXModel(vllm_config=vllm_config,
                                     prefix=maybe_prefix(prefix, "gpt_neox"))
258
        self.embed_out = ParallelLMHead(
259
            config.vocab_size,
260
            config.hidden_size,
261
            quant_config=quant_config,
262
        )
263
264
        if self.config.tie_word_embeddings:
            self.embed_out.weight = self.gpt_neox.embed_in.weight
265
        self.logits_processor = LogitsProcessor(config.vocab_size)
Joe Runde's avatar
Joe Runde committed
266
        self.sampler = get_sampler()
267
268
        self.make_empty_intermediate_tensors = (
            self.gpt_neox.make_empty_intermediate_tensors)
gaoqiong's avatar
gaoqiong committed
269
270
271
272
273
274
275
276
        
        self.w8a8_strategy=int(os.getenv('W8A8_SUPPORT_METHODS', '1'))
        self.quant_method = None
        if quant_config is not None:
            self.quant_method=quant_config.get_name()
            self.quant_config=quant_config
        self.tritonsingleton= W8a8GetCacheJSON()
        
277
278
279
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.gpt_neox.get_input_embeddings(input_ids)

280
281
    def forward(
        self,
282
283
        input_ids: torch.Tensor,
        positions: torch.Tensor,
284
        intermediate_tensors: Optional[IntermediateTensors] = None,
285
        inputs_embeds: Optional[torch.Tensor] = None,
286
    ) -> Union[torch.Tensor, IntermediateTensors]:
287
288
        hidden_states = self.gpt_neox(input_ids, positions,
                                      intermediate_tensors, inputs_embeds)
289
290
        return hidden_states

291
292
293
294
295
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
296
        logits = self.logits_processor(self.embed_out, hidden_states,
297
298
299
                                       sampling_metadata)
        return logits

300
301
    def sample(
        self,
302
        logits: torch.Tensor,
303
        sampling_metadata: SamplingMetadata,
304
    ) -> Optional[SamplerOutput]:
305
        next_tokens = self.sampler(logits, sampling_metadata)
306
307
        return next_tokens

308
309
    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
310
        params_dict = dict(self.named_parameters())
311
        loaded_params: Set[str] = set()
312
        for name, loaded_weight in weights:
313
            if ("attention.bias" in name or "attention.masked_bias" in name
314
                    or "rotary_emb.inv_freq" in name):
315
                continue
316
317
318
319
320
            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
321
322
            if is_pp_missing_parameter(name, self):
                continue
323
324
            param = params_dict[name]

325
            if "query_key_value" in name:
326
327
328
                # 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).
329
                # Thus, we need weight conversion.
330
                output_dim = getattr(param, "output_dim", None)
331
                num_heads = self.config.num_attention_heads
332
333
334
335
336
337
338
339
340
341
342
343
                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)
344
            loaded_params.add(name)
gaoqiong's avatar
gaoqiong committed
345
346
        #当为triton支持推理的时候不能进行处理
        if self.quant_method == "compressed_tensors":
347
            os.environ['LM_NN'] = '0'
gaoqiong's avatar
gaoqiong committed
348
349
350
351
352
353
354
355
356
            lay_key_words = [
                "attention.query_key_value.weight",
                "attention.dense.weight",
                "mlp.dense_h_to_4h.weight",
                "mlp.dense_4h_to_h.weight",
            ]
            combined_words = "|".join(lay_key_words)
            weight_shapes=[]
            all_json={}
357
            matched_key_words=set() 
gaoqiong's avatar
gaoqiong committed
358
359
360
361
362
363
364
365
366
367
368
369
370
371
            
            for layername, weight in params_dict.items():  
                matches = re.findall(combined_words, layername)
                if matches and "scale" not in layername:
                    weight_data =params_dict[layername]
                    n=weight_data.shape[0]
                    k=weight_data.shape[1]
       
                    #rocblas和cutlass目前都需要weight做处理,但是triton不用
                    if self.w8a8_strategy!=1:
                        _weight=weight_data.T.contiguous().reshape(n,-1)
                        weight_data.data.copy_(_weight)  
                    
                    #下面是针对模型记录模型出现k和n值 
372
373
374
                    elif len(matched_key_words) < 4 and matches[0] not in matched_key_words:
                        matched_key_words.add(matches[0])
  
gaoqiong's avatar
gaoqiong committed
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
                        weight_shapes.append({n,k})
                
                        json_file=self.tritonsingleton.get_w8a8json_name(n,k)
                        configs_dict=self.tritonsingleton.get_triton_cache(json_file,n,k)
                        if configs_dict:
                            all_json.update(configs_dict)
                                            
            if self.w8a8_strategy==1:
                self.tritonsingleton.triton_json_dict.append(all_json)
                #找到的所有config都进行一次warmup
                for key, value in all_json.items():
                    m=int(key.split('_')[0])
                    n=int(key.split('_')[1])
                    k=int(key.split('_')[2])
                    ops.triton_int8_gemm_helper(m=m,n=n,k=k,per_token_act_quant=True,per_out_channel_weight_quant=True,use_bias=False,best_config=value)
                   

392
        return loaded_params