baichuan.py 16.5 KB
Newer Older
codethazine's avatar
codethazine committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
# coding=utf-8
# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
#
# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
# and OPT implementations in this library. It has been modified from its
# original forms to accommodate minor architectural differences compared
# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
#
# 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
20
"""Inference-only BaiChuan model compatible with HuggingFace weights."""
21
import math
22
from typing import Iterable, List, Optional, Tuple
codethazine's avatar
codethazine committed
23
24

import torch
25
from torch import nn
26
from transformers import PretrainedConfig
codethazine's avatar
codethazine committed
27

28
from vllm.attention import Attention, AttentionMetadata
29
from vllm.config import CacheConfig, LoRAConfig
30
31
from vllm.distributed import (get_tensor_model_parallel_rank,
                              get_tensor_model_parallel_world_size)
32
33
from vllm.model_executor.layers.activation import SiluAndMul
from vllm.model_executor.layers.layernorm import RMSNorm
34
from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
35
36
                                               QKVParallelLinear,
                                               RowParallelLinear)
37
from vllm.model_executor.layers.logits_processor import LogitsProcessor
38
39
from vllm.model_executor.layers.quantization.base_config import (
    QuantizationConfig)
40
from vllm.model_executor.layers.rotary_embedding import get_rope
41
from vllm.model_executor.layers.sampler import Sampler, SamplerOutput
42
from vllm.model_executor.layers.vocab_parallel_embedding import (
43
    ParallelLMHead, VocabParallelEmbedding)
44
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
45
from vllm.model_executor.sampling_metadata import SamplingMetadata
46
from vllm.sequence import IntermediateTensors
codethazine's avatar
codethazine committed
47

48
49
from .interfaces import SupportsLoRA

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
78
79
80
81
82

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(start=1,
                                    end=1 + 2 * num_remaining_heads,
                                    step=2,
                                    dtype=torch.int32)
        slopes = torch.cat(
            [slopes, torch.pow(extra_base, extra_powers)], dim=0)
    return slopes


class BaiChuanMLP(nn.Module):

    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
        hidden_act: str,
83
        quant_config: Optional[QuantizationConfig] = None,
84
85
86
87
88
    ):
        super().__init__()
        self.gate_up_proj = MergedColumnParallelLinear(
            hidden_size, [intermediate_size] * 2,
            bias=False,
89
            quant_config=quant_config)
90
91
92
        self.down_proj = RowParallelLinear(intermediate_size,
                                           hidden_size,
                                           bias=False,
93
                                           quant_config=quant_config)
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
        if hidden_act != "silu":
            raise ValueError(f"Unsupported activation: {hidden_act}. "
                             "Only silu is supported for now.")
        self.act_fn = SiluAndMul()

    def forward(self, x):
        gate_up, _ = self.gate_up_proj(x)
        x = self.act_fn(gate_up)
        x, _ = self.down_proj(x)
        return x


class BaiChuanAttention(nn.Module):
    """Multi-headed attention from 'Attention Is All You Need' paper"""

    def __init__(
        self,
        hidden_size: int,
        num_heads: int,
        position_embedding: str,
        rope_theta: float = 10000,
        max_position_embeddings: int = 8192,
116
        cache_config: Optional[CacheConfig] = None,
117
        quant_config: Optional[QuantizationConfig] = None,
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
    ):
        super().__init__()
        self.hidden_size = hidden_size
        tensor_model_parallel_world_size = get_tensor_model_parallel_world_size(
        )
        self.total_num_heads = num_heads
        assert self.total_num_heads % tensor_model_parallel_world_size == 0
        self.num_heads = (self.total_num_heads //
                          tensor_model_parallel_world_size)
        self.head_dim = hidden_size // self.total_num_heads
        self.postion_embedding = position_embedding
        self.rope_theta = rope_theta
        self.max_position_embeddings = max_position_embeddings

        # pylint: disable=invalid-name
        self.W_pack = QKVParallelLinear(
            hidden_size,
            self.head_dim,
            self.total_num_heads,
            self.total_num_heads,
            bias=False,
139
            quant_config=quant_config,
140
141
142
143
144
        )
        self.o_proj = RowParallelLinear(
            self.total_num_heads * self.head_dim,
            hidden_size,
            bias=False,
145
            quant_config=quant_config,
146
147
148
149
150
151
152
153
154
155
        )
        # Create the alibi slopes and slice them.
        if self.postion_embedding == "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)
            alibi_slopes = alibi_slopes[head_start:head_end].tolist()

            scaling = self.head_dim**-0.5
156
157
158
            self.attn = Attention(self.num_heads,
                                  self.head_dim,
                                  scaling,
159
160
                                  alibi_slopes=alibi_slopes,
                                  quant_config=quant_config)
161
162
163
164
165
166
167
168
        else:
            self.rotary_emb = get_rope(
                self.head_dim,
                rotary_dim=self.head_dim,
                max_position=self.max_position_embeddings,
                base=self.rope_theta,
            )
            self.scaling = self.head_dim**-0.5
169
170
171
            self.attn = Attention(self.num_heads,
                                  self.head_dim,
                                  self.scaling,
172
173
                                  cache_config=cache_config,
                                  quant_config=quant_config)
174
175
176
177
178

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
179
180
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
181
182
183
184
185
    ) -> torch.Tensor:
        qkv, _ = self.W_pack(hidden_states)
        q, k, v = qkv.chunk(chunks=3, dim=-1)
        if self.postion_embedding != "ALIBI":
            q, k = self.rotary_emb(positions, q, k)
186
        attn_output = self.attn(q, k, v, kv_cache, attn_metadata)
187
188
189
190
191
192
193
        output, _ = self.o_proj(attn_output)
        return output


class BaiChuanDecoderLayer(nn.Module):

    def __init__(self,
194
                 config: PretrainedConfig,
195
                 position_embedding: str,
196
                 cache_config: Optional[CacheConfig] = None,
197
                 quant_config: Optional[QuantizationConfig] = None):
198
199
200
201
202
203
204
205
206
207
208
        super().__init__()
        self.hidden_size = config.hidden_size
        rope_theta = getattr(config, "rope_theta", 10000)
        max_position_embeddings = getattr(config, "max_position_embeddings",
                                          8192)
        self.self_attn = BaiChuanAttention(
            hidden_size=self.hidden_size,
            num_heads=config.num_attention_heads,
            position_embedding=position_embedding,
            rope_theta=rope_theta,
            max_position_embeddings=max_position_embeddings,
209
            cache_config=cache_config,
210
            quant_config=quant_config,
211
212
213
214
215
        )
        self.mlp = BaiChuanMLP(
            hidden_size=self.hidden_size,
            intermediate_size=config.intermediate_size,
            hidden_act=config.hidden_act,
216
            quant_config=quant_config,
217
218
219
220
221
222
223
224
225
226
        )
        self.input_layernorm = RMSNorm(config.hidden_size,
                                       eps=config.rms_norm_eps)
        self.post_attention_layernorm = RMSNorm(config.hidden_size,
                                                eps=config.rms_norm_eps)

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
227
228
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
229
230
231
232
233
234
235
236
237
238
239
240
241
        residual: Optional[torch.Tensor],
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        # Self Attention
        if residual is None:
            residual = hidden_states
            hidden_states = self.input_layernorm(hidden_states)
        else:
            hidden_states, residual = self.input_layernorm(
                hidden_states, residual)
        hidden_states = self.self_attn(
            positions=positions,
            hidden_states=hidden_states,
            kv_cache=kv_cache,
242
            attn_metadata=attn_metadata,
243
244
245
246
247
248
249
250
251
252
253
254
        )

        # Fully Connected
        hidden_states, residual = self.post_attention_layernorm(
            hidden_states, residual)
        hidden_states = self.mlp(hidden_states)
        return hidden_states, residual


class BaiChuanModel(nn.Module):

    def __init__(self,
255
                 config: PretrainedConfig,
256
                 position_embedding: str,
257
                 cache_config: Optional[CacheConfig] = None,
258
                 quant_config: Optional[QuantizationConfig] = None):
259
260
261
262
263
264
265
266
267
268
        super().__init__()
        self.config = config
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size

        self.embed_tokens = VocabParallelEmbedding(
            config.vocab_size,
            config.hidden_size,
        )
        self.layers = nn.ModuleList([
269
270
            BaiChuanDecoderLayer(config, position_embedding, cache_config,
                                 quant_config)
271
272
273
274
275
276
277
278
            for _ in range(config.num_hidden_layers)
        ])
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
279
280
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
281
282
283
284
285
286
287
288
289
    ) -> torch.Tensor:
        hidden_states = self.embed_tokens(input_ids)
        residual = None
        for i in range(len(self.layers)):
            layer = self.layers[i]
            hidden_states, residual = layer(
                positions,
                hidden_states,
                kv_caches[i],
290
                attn_metadata,
291
292
293
294
295
296
                residual,
            )
        hidden_states, _ = self.norm(hidden_states, residual)
        return hidden_states


297
class BaiChuanBaseForCausalLM(nn.Module, SupportsLoRA):
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
    packed_modules_mapping = {
        "W_pack": ["W_pack"],
        "gate_up_proj": [
            "gate_proj",
            "up_proj",
        ],
    }
    # LoRA specific attributes
    supported_lora_modules = [
        "W_pack",
        "o_proj",
        "gate_up_proj",
        "down_proj",
    ]
    embedding_modules = {}
    embedding_padding_modules = []
314

315
316
    def __init__(
        self,
317
        config: PretrainedConfig,
318
        position_embedding: str,
319
        cache_config: Optional[CacheConfig] = None,
320
        quant_config: Optional[QuantizationConfig] = None,
321
322
        lora_config: Optional[LoRAConfig] = None,
    ):
323
        super().__init__()
324

325
        self.config = config
326
327
        self.lora_config = lora_config

328
        self.quant_config = quant_config
329
330
        self.model = BaiChuanModel(config, position_embedding, cache_config,
                                   quant_config)
331
332
333
        self.lm_head = ParallelLMHead(config.vocab_size,
                                      config.hidden_size,
                                      quant_config=quant_config)
334
335
        if self.config.tie_word_embeddings:
            self.lm_head.weight = self.model.embed_tokens.weight
336
337
        self.logits_processor = LogitsProcessor(config.vocab_size)
        self.sampler = Sampler()
338
339
340
341
342

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
343
344
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
345
        intermediate_tensors: Optional[IntermediateTensors] = None,
346
347
    ) -> torch.Tensor:
        hidden_states = self.model(input_ids, positions, kv_caches,
348
                                   attn_metadata)
349
350
        return hidden_states

351
352
353
354
355
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
356
        logits = self.logits_processor(self.lm_head, hidden_states,
357
358
359
                                       sampling_metadata)
        return logits

360
361
    def sample(
        self,
362
        logits: torch.Tensor,
363
364
        sampling_metadata: SamplingMetadata,
    ) -> Optional[SamplerOutput]:
365
        next_tokens = self.sampler(logits, sampling_metadata)
366
        return next_tokens
codethazine's avatar
codethazine committed
367

368
    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
369
370
371
372
373
374
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
            ("gate_up_proj", "gate_proj", 0),
            ("gate_up_proj", "up_proj", 1),
        ]
        params_dict = dict(self.named_parameters())
375
        for name, loaded_weight in weights:
codethazine's avatar
codethazine committed
376
377
            if "rotary_emb.inv_freq" in name:
                continue
378
            if name == "lm_head.weight":
379
380
                # Unlike Baichuan, Baichuan2 normalizes the head weights.
                # Refer to:
381
382
383
384
385
386
387
388
389
                # https://huggingface.co/baichuan-inc/Baichuan2-7B-Chat/blob/84603cde5ebffb6084e476cfaeceaf0b8b91fe54/modeling_baichuan.py#L508
                # Distinguish between Baichuan and Baichuan2 by checking the
                # vocab size. This is suggested by
                # https://github.com/vllm-project/vllm/pull/1022#discussion_r1325652704
                is_baichuan2 = self.config.vocab_size == 125696
                if is_baichuan2:
                    loaded_weight = torch.nn.functional.normalize(
                        loaded_weight)

390
            for (param_name, weight_name, shard_id) in stacked_params_mapping:
codethazine's avatar
codethazine committed
391
392
                if weight_name not in name:
                    continue
CHU Tianxiang's avatar
CHU Tianxiang committed
393
394
395
396
397
                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
                param = params_dict[name]
398
399
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
codethazine's avatar
codethazine committed
400
                break
401
            else:
CHU Tianxiang's avatar
CHU Tianxiang committed
402
403
404
                # Skip loading extra bias for GPTQ models.
                if name.endswith(".bias") and name not in params_dict:
                    continue
405
406
407
408
                param = params_dict[name]
                weight_loader = getattr(param, "weight_loader",
                                        default_weight_loader)
                weight_loader(param, loaded_weight)
409
410


411
412
class BaichuanForCausalLM(BaiChuanBaseForCausalLM):
    """Baichuan 13B and Baichuan2 7B/13B."""
413

414
415
416
    def __init__(
        self,
        config,
417
        cache_config: Optional[CacheConfig] = None,
418
        quant_config: Optional[QuantizationConfig] = None,
419
420
        lora_config: Optional[LoRAConfig] = None,
    ):
421
        if config.hidden_size == 4096:  # baichuan2 7b
422
423
            super().__init__(config, "ROPE", cache_config, quant_config,
                             lora_config)
424
        else:  # baichuan 13b, baichuan2 13b
425
426
            super().__init__(config, "ALIBI", cache_config, quant_config,
                             lora_config)
427
428


429
430
class BaiChuanForCausalLM(BaiChuanBaseForCausalLM):
    """Baichuan 7B."""
431

432
433
434
    def __init__(
        self,
        config,
435
        cache_config: Optional[CacheConfig] = None,
436
        quant_config: Optional[QuantizationConfig] = None,
437
438
        lora_config: Optional[LoRAConfig] = None,
    ):
439
440
        super().__init__(config, "ROPE", cache_config, quant_config,
                         lora_config)