baichuan.py 18.6 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
# 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
19
"""Inference-only BaiChuan model compatible with HuggingFace weights."""
20
import math
21
from typing import Iterable, List, Optional, Set, Tuple, Union
codethazine's avatar
codethazine committed
22
23

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

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_rank,
31
                              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
from vllm.model_executor.layers.quantization import QuantizationConfig
39
from vllm.model_executor.layers.rotary_embedding import get_rope
Joe Runde's avatar
Joe Runde committed
40
from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
41
from vllm.model_executor.layers.vocab_parallel_embedding import (
42
    ParallelLMHead, VocabParallelEmbedding)
43
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
44
from vllm.model_executor.sampling_metadata import SamplingMetadata
45
from vllm.sequence import IntermediateTensors
codethazine's avatar
codethazine committed
46

47
48
49
from .interfaces import SupportsLoRA, SupportsPP
from .utils import (is_pp_missing_parameter,
                    make_empty_intermediate_tensors_factory, make_layers)
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
83

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,
84
        quant_config: Optional[QuantizationConfig] = None,
85
86
87
88
89
    ):
        super().__init__()
        self.gate_up_proj = MergedColumnParallelLinear(
            hidden_size, [intermediate_size] * 2,
            bias=False,
90
            quant_config=quant_config)
91
92
93
        self.down_proj = RowParallelLinear(intermediate_size,
                                           hidden_size,
                                           bias=False,
94
                                           quant_config=quant_config)
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
        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,
117
        cache_config: Optional[CacheConfig] = None,
118
        quant_config: Optional[QuantizationConfig] = None,
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
    ):
        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,
140
            quant_config=quant_config,
141
142
143
144
145
        )
        self.o_proj = RowParallelLinear(
            self.total_num_heads * self.head_dim,
            hidden_size,
            bias=False,
146
            quant_config=quant_config,
147
148
149
150
151
152
153
154
155
156
        )
        # 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
157
158
159
            self.attn = Attention(self.num_heads,
                                  self.head_dim,
                                  scaling,
160
161
                                  alibi_slopes=alibi_slopes,
                                  quant_config=quant_config)
162
163
164
165
166
167
168
169
        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
170
171
172
            self.attn = Attention(self.num_heads,
                                  self.head_dim,
                                  self.scaling,
173
174
                                  cache_config=cache_config,
                                  quant_config=quant_config)
175
176
177
178
179

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
180
181
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
182
183
184
185
186
    ) -> 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)
187
        attn_output = self.attn(q, k, v, kv_cache, attn_metadata)
188
189
190
191
192
193
194
        output, _ = self.o_proj(attn_output)
        return output


class BaiChuanDecoderLayer(nn.Module):

    def __init__(self,
195
                 config: PretrainedConfig,
196
                 position_embedding: str,
197
                 cache_config: Optional[CacheConfig] = None,
198
                 quant_config: Optional[QuantizationConfig] = None):
199
200
201
202
203
204
205
206
207
208
209
        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,
210
            cache_config=cache_config,
211
            quant_config=quant_config,
212
213
214
215
216
        )
        self.mlp = BaiChuanMLP(
            hidden_size=self.hidden_size,
            intermediate_size=config.intermediate_size,
            hidden_act=config.hidden_act,
217
            quant_config=quant_config,
218
219
220
221
222
223
224
225
226
227
        )
        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,
228
229
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
230
231
232
233
234
235
236
237
238
239
240
241
242
        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,
243
            attn_metadata=attn_metadata,
244
245
246
247
248
249
250
251
252
        )

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


253
@support_torch_compile
254
255
class BaiChuanModel(nn.Module):

256
257
258
259
260
261
    def __init__(
        self,
        vllm_config: VllmConfig,
        prefix: str = "",
        position_embedding: str = "ROPE",
    ) -> None:
262
        super().__init__()
263
264
265
266
267

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

268
269
270
271
272
273
274
275
        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,
        )
276
277
278
279
280
281
        self.start_layer, self.end_layer, self.layers = make_layers(
            config.num_hidden_layers,
            lambda prefix: BaiChuanDecoderLayer(config, position_embedding,
                                                cache_config, quant_config),
            prefix=f"{prefix}.layers",
        )
282
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
283
284
285
        self.make_empty_intermediate_tensors = (
            make_empty_intermediate_tensors_factory(
                ["hidden_states", "residual"], config.hidden_size))
286

287
288
289
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embed_tokens(input_ids)

290
291
292
293
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
294
295
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
296
        intermediate_tensors: Optional[IntermediateTensors],
297
        inputs_embeds: Optional[torch.Tensor] = None,
298
299
    ) -> Union[torch.Tensor, IntermediateTensors]:
        if get_pp_group().is_first_rank:
300
301
302
303
            if inputs_embeds is not None:
                hidden_states = inputs_embeds
            else:
                hidden_states = self.get_input_embeddings(input_ids)
304
305
306
307
308
309
            residual = None
        else:
            assert intermediate_tensors is not None
            hidden_states = intermediate_tensors["hidden_states"]
            residual = intermediate_tensors["residual"]
        for i in range(self.start_layer, self.end_layer):
310
311
312
313
            layer = self.layers[i]
            hidden_states, residual = layer(
                positions,
                hidden_states,
314
                kv_caches[i - self.start_layer],
315
                attn_metadata,
316
317
                residual,
            )
318
319
320
321
322
        if not get_pp_group().is_last_rank:
            return IntermediateTensors({
                "hidden_states": hidden_states,
                "residual": residual,
            })
323
324
325
326
        hidden_states, _ = self.norm(hidden_states, residual)
        return hidden_states


327
class BaiChuanBaseForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
    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 = []
344

345
346
    def __init__(
        self,
347
        *,
348
349
350
        vllm_config: VllmConfig,
        prefix: str = "",
        position_embedding: str = "ROPE",
351
    ):
352
        super().__init__()
353
354
355
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
        lora_config = vllm_config.lora_config
356
        self.config = config
357
358
        self.lora_config = lora_config

359
        self.quant_config = quant_config
360
361
362
        self.model = BaiChuanModel(vllm_config=vllm_config,
                                   prefix=prefix,
                                   position_embedding=position_embedding)
363
364
365
        self.lm_head = ParallelLMHead(config.vocab_size,
                                      config.hidden_size,
                                      quant_config=quant_config)
366
367
        if self.config.tie_word_embeddings:
            self.lm_head.weight = self.model.embed_tokens.weight
368
        self.logits_processor = LogitsProcessor(config.vocab_size)
Joe Runde's avatar
Joe Runde committed
369
        self.sampler = get_sampler()
370
371
        self.make_empty_intermediate_tensors = (
            self.model.make_empty_intermediate_tensors)
372

373
374
375
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.model.get_input_embeddings(input_ids)

376
377
378
379
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
380
381
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
382
        intermediate_tensors: Optional[IntermediateTensors] = None,
383
        inputs_embeds: Optional[torch.Tensor] = None,
384
    ) -> Union[torch.Tensor, IntermediateTensors]:
385
        hidden_states = self.model(input_ids, positions, kv_caches,
386
387
                                   attn_metadata, intermediate_tensors,
                                   inputs_embeds)
388
389
        return hidden_states

390
391
392
393
394
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
395
        logits = self.logits_processor(self.lm_head, hidden_states,
396
397
398
                                       sampling_metadata)
        return logits

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

407
408
    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
409
410
411
412
413
414
        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())
415
        loaded_params: Set[str] = set()
416
        for name, loaded_weight in weights:
codethazine's avatar
codethazine committed
417
418
            if "rotary_emb.inv_freq" in name:
                continue
419
            if name == "lm_head.weight":
420
421
                # Unlike Baichuan, Baichuan2 normalizes the head weights.
                # Refer to:
422
423
424
425
426
427
428
429
430
                # 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)

431
            for (param_name, weight_name, shard_id) in stacked_params_mapping:
codethazine's avatar
codethazine committed
432
433
                if weight_name not in name:
                    continue
CHU Tianxiang's avatar
CHU Tianxiang committed
434
435
436
437
                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
438
439
                if is_pp_missing_parameter(name, self):
                    continue
CHU Tianxiang's avatar
CHU Tianxiang committed
440
                param = params_dict[name]
441
442
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
codethazine's avatar
codethazine committed
443
                break
444
            else:
CHU Tianxiang's avatar
CHU Tianxiang committed
445
446
447
                # Skip loading extra bias for GPTQ models.
                if name.endswith(".bias") and name not in params_dict:
                    continue
448
449
                if is_pp_missing_parameter(name, self):
                    continue
450
451
452
453
                param = params_dict[name]
                weight_loader = getattr(param, "weight_loader",
                                        default_weight_loader)
                weight_loader(param, loaded_weight)
454
455
            loaded_params.add(name)
        return loaded_params
456
457


458
class BaichuanForCausalLM(BaiChuanBaseForCausalLM):
459
460
461
    """Baichuan 13B and Baichuan2 7B/13B.
    NOTE: the class name has a lower case 'c'.
    """
462

463
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
464
        config = vllm_config.model_config.hf_config
465
        if config.hidden_size == 4096:  # baichuan2 7b
466
467
468
            super().__init__(vllm_config=vllm_config,
                             prefix=prefix,
                             position_embedding="ROPE")
469
        else:  # baichuan 13b, baichuan2 13b
470
471
472
            super().__init__(vllm_config=vllm_config,
                             prefix=prefix,
                             position_embedding="ALIBI")
473
474


475
class BaiChuanForCausalLM(BaiChuanBaseForCausalLM):
476
477
478
    """Baichuan 7B.
    NOTE: the class name has an upper case 'C'.
    """
479

480
481
482
483
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        super().__init__(vllm_config=vllm_config,
                         prefix=prefix,
                         position_embedding="ROPE")