gpt_neox.py 16.4 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, List, Optional, Set, Tuple, Union
gaoqiong's avatar
gaoqiong committed
21
22
import os
import re
23
24
import torch
from torch import nn
25
26
from transformers import GPTNeoXConfig

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

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

class GPTNeoXAttention(nn.Module):

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

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

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

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


class GPTNeoXMLP(nn.Module):
121

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

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

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

    def forward(
        self,
170
        position_ids: torch.Tensor,
171
        hidden_states: torch.Tensor,
172
173
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
174
175
176
177
178
179
    ) -> torch.Tensor:
        attn_input = self.input_layernorm(hidden_states)
        attn_output = self.attention(
            position_ids=position_ids,
            hidden_states=attn_input,
            kv_cache=kv_cache,
180
            attn_metadata=attn_metadata,
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
        )

        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


200
@support_torch_compile
201
class GPTNeoXModel(nn.Module):
Woosuk Kwon's avatar
Woosuk Kwon committed
202

203
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
204
        super().__init__()
205
206
207
208
209

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

210
211
        self.config = config

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

228
229
230
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embed_in(input_ids)

231
232
    def forward(
        self,
233
234
        input_ids: torch.Tensor,
        position_ids: torch.Tensor,
235
236
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
237
        intermediate_tensors: Optional[IntermediateTensors],
238
        inputs_embeds: Optional[torch.Tensor] = None,
239
240
    ) -> Union[torch.Tensor, IntermediateTensors]:
        if get_pp_group().is_first_rank:
241
242
243
244
            if inputs_embeds is not None:
                hidden_states = inputs_embeds
            else:
                hidden_states = self.get_input_embeddings(input_ids)
245
246
247
        else:
            hidden_states = intermediate_tensors["hidden_states"]
        for i in range(self.start_layer, self.end_layer):
248
249
250
251
            layer = self.layers[i]
            hidden_states = layer(
                position_ids,
                hidden_states,
252
                kv_caches[i - self.start_layer],
253
                attn_metadata,
254
            )
255
256
        if not get_pp_group().is_last_rank:
            return IntermediateTensors({"hidden_states": hidden_states})
257
258
259
260
        hidden_states = self.final_layer_norm(hidden_states)
        return hidden_states


261
class GPTNeoXForCausalLM(nn.Module, SupportsPP):
262

263
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
264
        super().__init__()
265
266
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
267
        self.config = config
268
        self.quant_config = quant_config
269
270
        self.gpt_neox = GPTNeoXModel(vllm_config=vllm_config,
                                     prefix=maybe_prefix(prefix, "gpt_neox"))
271
        self.embed_out = ParallelLMHead(
272
            config.vocab_size,
273
            config.hidden_size,
274
            quant_config=quant_config,
275
        )
276
277
        if self.config.tie_word_embeddings:
            self.embed_out.weight = self.gpt_neox.embed_in.weight
278
        self.logits_processor = LogitsProcessor(config.vocab_size)
Joe Runde's avatar
Joe Runde committed
279
        self.sampler = get_sampler()
280
281
        self.make_empty_intermediate_tensors = (
            self.gpt_neox.make_empty_intermediate_tensors)
gaoqiong's avatar
gaoqiong committed
282
283
284
285
286
287
288
289
        
        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()
        
290
291
292
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.gpt_neox.get_input_embeddings(input_ids)

293
294
    def forward(
        self,
295
296
        input_ids: torch.Tensor,
        positions: torch.Tensor,
297
298
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
299
        intermediate_tensors: Optional[IntermediateTensors] = None,
300
        inputs_embeds: Optional[torch.Tensor] = None,
301
    ) -> Union[torch.Tensor, IntermediateTensors]:
302
        hidden_states = self.gpt_neox(input_ids, positions, kv_caches,
303
304
                                      attn_metadata, intermediate_tensors,
                                      inputs_embeds)
305
306
        return hidden_states

307
308
309
310
311
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
312
        logits = self.logits_processor(self.embed_out, hidden_states,
313
314
315
                                       sampling_metadata)
        return logits

316
317
    def sample(
        self,
318
        logits: torch.Tensor,
319
        sampling_metadata: SamplingMetadata,
320
    ) -> Optional[SamplerOutput]:
321
        next_tokens = self.sampler(logits, sampling_metadata)
322
323
        return next_tokens

324
325
    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
326
        params_dict = dict(self.named_parameters())
327
        loaded_params: Set[str] = set()
328
        for name, loaded_weight in weights:
329
            if ("attention.bias" in name or "attention.masked_bias" in name
330
                    or "rotary_emb.inv_freq" in name):
331
                continue
332
333
334
335
336
            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
337
338
            if is_pp_missing_parameter(name, self):
                continue
339
340
            param = params_dict[name]

341
            if "query_key_value" in name:
342
343
344
                # 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).
345
                # Thus, we need weight conversion.
346
                output_dim = getattr(param, "output_dim", None)
347
                num_heads = self.config.num_attention_heads
348
349
350
351
352
353
354
355
356
357
358
359
                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)
360
            loaded_params.add(name)
gaoqiong's avatar
gaoqiong committed
361
362
        #当为triton支持推理的时候不能进行处理
        if self.quant_method == "compressed_tensors":
363
            os.environ['LM_NN'] = '0'
gaoqiong's avatar
gaoqiong committed
364
365
366
367
368
369
370
371
372
            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={}
373
            matched_key_words=set() 
gaoqiong's avatar
gaoqiong committed
374
375
376
377
378
379
380
381
382
383
384
385
386
387
            
            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值 
388
389
390
                    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
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
                        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)
                   

408
        return loaded_params