falcon.py 17.7 KB
Newer Older
Zhuohan Li's avatar
Zhuohan Li committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
# coding=utf-8
# Adapted from
# https://github.com/huggingface/transformers/blob/a5cc30d72ae2dc19af534e4b35c986cc28db1275/src/transformers/models/falcon/modeling_falcon.py
# Copyright 2023 The vLLM team.
# Copyright 2023 the Falcon authors and 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.
"""PyTorch Falcon model."""

import math
22
from typing import List, Optional, Union
Zhuohan Li's avatar
Zhuohan Li committed
23
24
25
26
27
28

import torch
from torch import nn
from torch.nn import LayerNorm
from transformers import FalconConfig as HF_FalconConfig

29
from vllm.attention import Attention, AttentionMetadata
30
from vllm.model_executor.layers.activation import get_act_fn
31
32
33
34
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
                                               LinearMethodBase,
                                               QKVParallelLinear,
                                               RowParallelLinear)
35
from vllm.model_executor.layers.logits_processor import LogitsProcessor
36
from vllm.model_executor.layers.rotary_embedding import get_rope
Zhuohan Li's avatar
Zhuohan Li committed
37
from vllm.model_executor.layers.sampler import Sampler
38
from vllm.model_executor.layers.vocab_parallel_embedding import (
39
    VocabParallelEmbedding)
40
41
from vllm.model_executor.parallel_utils.communication_op import (
    tensor_model_parallel_all_reduce)
42
43
from vllm.model_executor.parallel_utils.parallel_state import (
    get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size)
44
from vllm.model_executor.sampling_metadata import SamplingMetadata
45
46
from vllm.model_executor.weight_utils import (default_weight_loader,
                                              hf_model_weights_iterator)
47
from vllm.sequence import SamplerOutput
Zhuohan Li's avatar
Zhuohan Li committed
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
from vllm.transformers_utils.configs import RWConfig

FalconConfig = Union[HF_FalconConfig, RWConfig]


def _get_alibi_slopes(total_num_heads: int) -> torch.Tensor:
    closest_power_of_2 = 2**math.floor(math.log2(total_num_heads))
    base = torch.tensor(2**(-(2**-(math.log2(closest_power_of_2) - 3))),
                        dtype=torch.float32)
    powers = torch.arange(1, 1 + closest_power_of_2, dtype=torch.int32)
    slopes = torch.pow(base, powers)

    if closest_power_of_2 != total_num_heads:
        extra_base = torch.tensor(
            2**(-(2**-(math.log2(2 * closest_power_of_2) - 3))),
            dtype=torch.float32)
        num_remaining_heads = min(closest_power_of_2,
                                  total_num_heads - closest_power_of_2)
        extra_powers = torch.arange(1,
                                    1 + 2 * num_remaining_heads,
                                    2,
                                    dtype=torch.int32)
        slopes = torch.cat(
            [slopes, torch.pow(extra_base, extra_powers)], dim=0)

    return slopes


class FalconAttention(nn.Module):

78
79
80
81
82
    def __init__(
        self,
        config: FalconConfig,
        linear_method: Optional[LinearMethodBase] = None,
    ):
Zhuohan Li's avatar
Zhuohan Li committed
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
        super().__init__()

        self.hidden_size = config.hidden_size
        tp_size = get_tensor_model_parallel_world_size()

        self.total_num_heads = config.num_attention_heads
        assert self.total_num_heads % tp_size == 0
        self.num_heads = self.total_num_heads // tp_size
        self.head_dim = self.hidden_size // self.total_num_heads
        assert self.head_dim * self.total_num_heads == self.hidden_size

        self.new_decoder_architecture = config.new_decoder_architecture
        self.multi_query = config.multi_query

        if self.new_decoder_architecture:
            self.total_num_kv_heads = config.num_kv_heads
        elif self.multi_query:
            self.total_num_kv_heads = 1
        else:
            self.total_num_kv_heads = self.total_num_heads
103
104
105
106
107
108
109
110
111
        if self.total_num_kv_heads >= tp_size:
            # Number of KV heads is greater than TP size, so we partition
            # the KV heads across multiple tensor parallel GPUs.
            assert self.total_num_kv_heads % tp_size == 0
        else:
            # Number of KV heads is less than TP size, so we replicate
            # the KV heads across multiple tensor parallel GPUs.
            assert tp_size % self.total_num_kv_heads == 0
        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
Zhuohan Li's avatar
Zhuohan Li committed
112

113
114
115
116
117
118
119
120
121
        self.query_key_value = QKVParallelLinear(
            self.hidden_size,
            self.head_dim,
            self.total_num_heads,
            self.total_num_kv_heads,
            bias=config.bias,
            skip_bias_add=True,
            linear_method=linear_method,
        )
Zhuohan Li's avatar
Zhuohan Li committed
122
123
124
125
126
127
128
129
130
131
132
133
        self.q_size = self.num_heads * self.head_dim
        self.kv_size = self.num_kv_heads * self.head_dim

        # Layer-wise attention scaling
        self.inv_norm_factor = 1.0 / math.sqrt(self.head_dim)
        self.reduce_row_parallel_results = not (config.new_decoder_architecture
                                                or config.parallel_attn)
        self.dense = RowParallelLinear(
            self.hidden_size,
            self.hidden_size,
            bias=config.bias,
            skip_bias_add=True,
134
            linear_method=linear_method,
Zhuohan Li's avatar
Zhuohan Li committed
135
136
137
138
139
140
141
142
            reduce_results=self.reduce_row_parallel_results)

        self.use_rotary = config.rotary
        self.use_alibi = config.alibi
        assert not (self.use_rotary and self.use_alibi), (
            "Rotary and alibi are mutually exclusive.")

        if self.use_rotary:
143
144
145
            rope_theta = getattr(config, "rope_theta", 10000)
            max_position_embeddings = getattr(config,
                                              "max_position_embeddings", 8192)
Woosuk Kwon's avatar
Woosuk Kwon committed
146
            self.rotary_emb = get_rope(
147
148
                self.head_dim,
                rotary_dim=self.head_dim,
Woosuk Kwon's avatar
Woosuk Kwon committed
149
150
151
                max_position=max_position_embeddings,
                base=rope_theta,
            )
152
153
154
155
            self.attn = Attention(self.num_heads,
                                  self.head_dim,
                                  self.inv_norm_factor,
                                  num_kv_heads=self.num_kv_heads)
Zhuohan Li's avatar
Zhuohan Li committed
156
157
158
159
160
161
162
        elif self.use_alibi:
            tp_rank = get_tensor_model_parallel_rank()
            head_start = tp_rank * self.num_heads
            head_end = (tp_rank + 1) * self.num_heads
            alibi_slopes = (_get_alibi_slopes(self.total_num_heads) *
                            self.inv_norm_factor)
            alibi_slopes = alibi_slopes[head_start:head_end].tolist()
163
164
165
166
167
            self.attn = Attention(self.num_heads,
                                  self.head_dim,
                                  self.inv_norm_factor,
                                  num_kv_heads=self.num_kv_heads,
                                  alibi_slopes=alibi_slopes)
Zhuohan Li's avatar
Zhuohan Li committed
168
        else:
169
170
171
172
            self.attn = Attention(self.num_heads,
                                  self.head_dim,
                                  scale=self.inv_norm_factor,
                                  num_kv_heads=self.num_kv_heads)
Zhuohan Li's avatar
Zhuohan Li committed
173
174
175
176
177

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
178
179
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
Zhuohan Li's avatar
Zhuohan Li committed
180
    ) -> torch.Tensor:
181
182
183
184
        qkv, bias = self.query_key_value(hidden_states)
        if bias is not None:
            qkv += bias
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
Zhuohan Li's avatar
Zhuohan Li committed
185
        if self.use_rotary:
Woosuk Kwon's avatar
Woosuk Kwon committed
186
            q, k = self.rotary_emb(positions, q, k)
187
        attn_output = self.attn(q, k, v, kv_cache, attn_metadata)
Zhuohan Li's avatar
Zhuohan Li committed
188
189
190
191
192
193
        attn_output, bias = self.dense(attn_output)
        return attn_output, bias


class FalconMLP(nn.Module):

194
195
196
197
198
    def __init__(
        self,
        config: FalconConfig,
        linear_method: Optional[LinearMethodBase] = None,
    ):
Zhuohan Li's avatar
Zhuohan Li committed
199
200
201
202
203
204
        super().__init__()
        hidden_size = config.hidden_size

        self.dense_h_to_4h = ColumnParallelLinear(hidden_size,
                                                  4 * hidden_size,
                                                  bias=config.bias,
205
206
                                                  skip_bias_add=True,
                                                  linear_method=linear_method)
207
208
        quant_config = getattr(linear_method, "quant_config", None)
        self.act = get_act_fn("gelu", quant_config, 4 * hidden_size)
Zhuohan Li's avatar
Zhuohan Li committed
209
210
211
212
213
214
215
        self.reduce_row_parallel_results = not (config.new_decoder_architecture
                                                or config.parallel_attn)
        self.dense_4h_to_h = RowParallelLinear(
            4 * hidden_size,
            hidden_size,
            bias=config.bias,
            skip_bias_add=True,
216
217
            reduce_results=self.reduce_row_parallel_results,
            linear_method=linear_method)
Zhuohan Li's avatar
Zhuohan Li committed
218
219
220
221
222
223
224
225
226
227
228
229
230

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # NOTE(zhuohan): Following huggingface, we do not fuse bias add here.
        x, bias = self.dense_h_to_4h(x)
        if bias is not None:
            x += bias
        x = self.act(x)
        x, bias = self.dense_4h_to_h(x)
        return x, bias


class FalconDecoderLayer(nn.Module):

231
232
233
234
235
    def __init__(
        self,
        config: FalconConfig,
        linear_method: Optional[LinearMethodBase] = None,
    ):
Zhuohan Li's avatar
Zhuohan Li committed
236
237
238
        super().__init__()
        hidden_size = config.hidden_size
        self.num_heads = config.num_attention_heads
239
240
        self.self_attention = FalconAttention(config, linear_method)
        self.mlp = FalconMLP(config, linear_method)
Zhuohan Li's avatar
Zhuohan Li committed
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
        self.config = config

        if config.new_decoder_architecture:
            # The layer norm before self-attention
            self.ln_attn = LayerNorm(hidden_size,
                                     eps=config.layer_norm_epsilon)
            # The layer norm before the MLP
            self.ln_mlp = LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
        else:
            self.input_layernorm = LayerNorm(hidden_size,
                                             eps=config.layer_norm_epsilon)
            if not config.parallel_attn:
                self.post_attention_layernorm = LayerNorm(
                    hidden_size, eps=config.layer_norm_epsilon)

        self.reduce_row_parallel_results = not (config.new_decoder_architecture
                                                or config.parallel_attn)

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
263
264
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
265
    ) -> torch.Tensor:
Zhuohan Li's avatar
Zhuohan Li committed
266
267
268
269
270
271
272
273
274
275
276
277
278
        residual = hidden_states

        if self.config.new_decoder_architecture:
            attention_layernorm_out = self.ln_attn(hidden_states)
            mlp_layernorm_out = self.ln_mlp(hidden_states)
        else:
            attention_layernorm_out = self.input_layernorm(hidden_states)

        # Self attention.
        attention_output, attention_bias = self.self_attention(
            positions=positions,
            hidden_states=attention_layernorm_out,
            kv_cache=kv_cache,
279
            attn_metadata=attn_metadata,
Zhuohan Li's avatar
Zhuohan Li committed
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
        )
        if self.reduce_row_parallel_results and attention_bias is not None:
            attention_output += attention_bias

        if not self.config.new_decoder_architecture:
            if self.config.parallel_attn:
                mlp_layernorm_out = attention_layernorm_out
            else:
                residual += attention_output
                mlp_layernorm_out = self.post_attention_layernorm(residual)

        # MLP.
        mlp_output, mlp_bias = self.mlp(mlp_layernorm_out)
        if self.reduce_row_parallel_results and mlp_bias is not None:
            mlp_output += mlp_bias

        if not self.reduce_row_parallel_results:
            # When MLP and Attention layers are parallel, we can use
            # only one all-reduce operator to reduce the results from
            # both MLP and Attention layers.
            mlp_output += attention_output
301
            mlp_output = tensor_model_parallel_all_reduce(mlp_output)
Zhuohan Li's avatar
Zhuohan Li committed
302
303
304
305
306
307
308
309
310
311
312
            if attention_bias is not None:
                mlp_output += attention_bias
            if mlp_bias is not None:
                mlp_output += mlp_bias

        output = mlp_output + residual
        return output


class FalconModel(nn.Module):

313
314
315
316
317
    def __init__(
        self,
        config: FalconConfig,
        linear_method: Optional[LinearMethodBase] = None,
    ):
Zhuohan Li's avatar
Zhuohan Li committed
318
319
320
321
322
323
324
325
        super().__init__()
        self.config = config
        self.embed_dim = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.use_alibi = config.alibi

        # Embedding + LN Embedding
        self.word_embeddings = VocabParallelEmbedding(
326
327
328
            config.vocab_size,
            self.embed_dim,
        )
Zhuohan Li's avatar
Zhuohan Li committed
329
330
331

        # Transformer blocks
        self.h = nn.ModuleList([
332
333
            FalconDecoderLayer(config, linear_method)
            for _ in range(config.num_hidden_layers)
Zhuohan Li's avatar
Zhuohan Li committed
334
335
336
337
338
339
340
341
342
        ])

        # Final Layer Norm
        self.ln_f = LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)

    def forward(
        self,
        input_ids: torch.LongTensor,
        positions: torch.Tensor,
343
344
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
Zhuohan Li's avatar
Zhuohan Li committed
345
346
347
348
349
350
351
352
    ) -> torch.Tensor:
        hidden_states = self.word_embeddings(input_ids)
        for i in range(len(self.h)):
            layer = self.h[i]
            hidden_states = layer(
                positions,
                hidden_states,
                kv_caches[i],
353
                attn_metadata,
Zhuohan Li's avatar
Zhuohan Li committed
354
355
356
357
358
359
360
            )
        hidden_states = self.ln_f(hidden_states)
        return hidden_states


class FalconForCausalLM(nn.Module):

361
362
363
364
365
    def __init__(
        self,
        config: FalconConfig,
        linear_method: Optional[LinearMethodBase] = None,
    ):
Zhuohan Li's avatar
Zhuohan Li committed
366
367
        super().__init__()
        self.config = config
368
369
        self.linear_method = linear_method
        self.transformer = FalconModel(config, linear_method)
370
        self.lm_head_weight = self.transformer.word_embeddings.weight
371
372
        self.logits_processor = LogitsProcessor(config.vocab_size)
        self.sampler = Sampler()
Zhuohan Li's avatar
Zhuohan Li committed
373
374
375
376
377

    def forward(
        self,
        input_ids: torch.LongTensor,
        positions: torch.Tensor,
378
379
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
380
    ) -> torch.Tensor:
Zhuohan Li's avatar
Zhuohan Li committed
381
382
383
384
        hidden_states = self.transformer(
            input_ids,
            positions,
            kv_caches,
385
            attn_metadata,
Zhuohan Li's avatar
Zhuohan Li committed
386
        )
387
        return hidden_states
Zhuohan Li's avatar
Zhuohan Li committed
388

389
390
    def compute_logits(self, hidden_states: torch.Tensor,
                       sampling_metadata: SamplingMetadata) -> torch.Tensor:
391
        logits = self.logits_processor(self.lm_head_weight, hidden_states,
392
393
394
                                       sampling_metadata)
        return logits

395
396
    def sample(
        self,
397
        logits: torch.Tensor,
398
        sampling_metadata: SamplingMetadata,
399
    ) -> Optional[SamplerOutput]:
400
        next_tokens = self.sampler(logits, sampling_metadata)
Zhuohan Li's avatar
Zhuohan Li committed
401
402
403
404
405
        return next_tokens

    def load_weights(self,
                     model_name_or_path: str,
                     cache_dir: Optional[str] = None,
Jasmond L's avatar
Jasmond L committed
406
407
                     load_format: str = "auto",
                     revision: Optional[str] = None):
Zhuohan Li's avatar
Zhuohan Li committed
408
409
410
411
412
413
414
415
        total_num_heads = self.config.num_attention_heads
        if self.config.new_decoder_architecture:
            total_num_kv_heads = self.config.num_kv_heads
        elif self.config.multi_query:
            total_num_kv_heads = 1
        else:
            total_num_kv_heads = total_num_heads
        num_query_heads_per_kv_head = total_num_heads // total_num_kv_heads
416
        params_dict = dict(self.named_parameters(remove_duplicate=False))
Zhuohan Li's avatar
Zhuohan Li committed
417
        for name, loaded_weight in hf_model_weights_iterator(
Jasmond L's avatar
Jasmond L committed
418
                model_name_or_path, cache_dir, load_format, revision):
419
420
421
            if name == "lm_head.weight":
                # Falcon uses tied embeddings.
                continue
CHU Tianxiang's avatar
CHU Tianxiang committed
422
423
424
            # Skip loading extra bias for GPTQ models.
            if name.endswith(".bias") and name not in params_dict:
                continue
425
            param = params_dict[name]
Zhuohan Li's avatar
Zhuohan Li committed
426
            if "query_key_value" in name:
427
428
                output_dim = getattr(param, "output_dim", None)
                loaded_weight_shape = loaded_weight.shape
CHU Tianxiang's avatar
CHU Tianxiang committed
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
                if output_dim is not None:
                    loaded_weight = loaded_weight.view(
                        loaded_weight_shape[:output_dim] +
                        (total_num_kv_heads, num_query_heads_per_kv_head + 2,
                         -1) + loaded_weight_shape[output_dim + 1:])
                    wq = loaded_weight.narrow(
                        output_dim + 1, 0,
                        num_query_heads_per_kv_head).reshape(
                            *loaded_weight_shape[:output_dim], -1,
                            *loaded_weight_shape[output_dim + 1:])
                    wk = loaded_weight.narrow(
                        output_dim + 1, num_query_heads_per_kv_head,
                        1).reshape(*loaded_weight_shape[:output_dim], -1,
                                   *loaded_weight_shape[output_dim + 1:])
                    wv = loaded_weight.narrow(
                        output_dim + 1, num_query_heads_per_kv_head + 1,
                        1).reshape(*loaded_weight_shape[:output_dim], -1,
                                   *loaded_weight_shape[output_dim + 1:])
                    loaded_weight = torch.cat([wq, wk, wv], dim=output_dim)
448
449
450
451

            weight_loader = getattr(param, "weight_loader",
                                    default_weight_loader)
            weight_loader(param, loaded_weight)