baichuan.py 17.5 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

codethazine's avatar
codethazine committed
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
# 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
22
"""Inference-only BaiChuan model compatible with HuggingFace weights."""
23

24
import math
25
from collections.abc import Iterable
26
from itertools import islice
codethazine's avatar
codethazine committed
27
28

import torch
29
from torch import nn
30
from transformers import PretrainedConfig
codethazine's avatar
codethazine committed
31

32
from vllm.attention import Attention
33
from vllm.compilation.decorators import support_torch_compile
34
from vllm.config import CacheConfig, VllmConfig
35
36
37
38
39
from vllm.distributed import (
    get_pp_group,
    get_tensor_model_parallel_rank,
    get_tensor_model_parallel_world_size,
)
40
41
from vllm.model_executor.layers.activation import SiluAndMul
from vllm.model_executor.layers.layernorm import RMSNorm
42
43
44
45
46
from vllm.model_executor.layers.linear import (
    MergedColumnParallelLinear,
    QKVParallelLinear,
    RowParallelLinear,
)
47
from vllm.model_executor.layers.logits_processor import LogitsProcessor
48
from vllm.model_executor.layers.quantization import QuantizationConfig
49
from vllm.model_executor.layers.rotary_embedding import get_rope
50
from vllm.model_executor.layers.vocab_parallel_embedding import (
51
52
53
    ParallelLMHead,
    VocabParallelEmbedding,
)
54
from vllm.model_executor.model_loader.weight_utils import (
55
56
57
    default_weight_loader,
    row_parallel_weight_loader,
)
58
from vllm.sequence import IntermediateTensors
codethazine's avatar
codethazine committed
59

60
from .interfaces import SupportsLoRA, SupportsPP, SupportsQuant
61
62
63
64
65
66
67
from .utils import (
    AutoWeightsLoader,
    is_pp_missing_parameter,
    make_empty_intermediate_tensors_factory,
    make_layers,
    maybe_prefix,
)
68

69
70

def _get_alibi_slopes(total_num_heads: int) -> torch.Tensor:
71
    closest_power_of_2 = 2 ** math.floor(math.log2(total_num_heads))
72
    base = torch.tensor(
73
        2 ** (-(2 ** -(math.log2(closest_power_of_2) - 3))),
74
75
76
77
78
79
80
        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(
81
            2 ** (-(2 ** -(math.log2(2 * closest_power_of_2) - 3))),
82
83
            dtype=torch.float32,
        )
84
85
86
87
88
89
90
        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)
91
92
93
94
95
96
97
98
99
    return slopes


class BaiChuanMLP(nn.Module):
    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
        hidden_act: str,
100
        quant_config: QuantizationConfig | None = None,
101
        prefix: str = "",
102
103
104
    ):
        super().__init__()
        self.gate_up_proj = MergedColumnParallelLinear(
105
106
107
108
109
            hidden_size,
            [intermediate_size] * 2,
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.gate_up_proj",
110
111
        )
        self.down_proj = RowParallelLinear(
112
113
114
115
116
            intermediate_size,
            hidden_size,
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.down_proj",
117
        )
118
        if hidden_act != "silu":
119
120
121
            raise ValueError(
                f"Unsupported activation: {hidden_act}. Only silu is supported for now."
            )
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
        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,
141
142
        cache_config: CacheConfig | None = None,
        quant_config: QuantizationConfig | None = None,
143
        prefix: str = "",
144
145
146
    ):
        super().__init__()
        self.hidden_size = hidden_size
147
        tensor_model_parallel_world_size = get_tensor_model_parallel_world_size()
148
149
        self.total_num_heads = num_heads
        assert self.total_num_heads % tensor_model_parallel_world_size == 0
150
        self.num_heads = self.total_num_heads // tensor_model_parallel_world_size
151
        self.head_dim = hidden_size // self.total_num_heads
152
        self.position_embedding = position_embedding
153
154
155
156
157
158
159
160
161
162
        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,
163
            quant_config=quant_config,
164
            prefix=f"{prefix}.W_pack",
165
166
167
168
169
        )
        self.o_proj = RowParallelLinear(
            self.total_num_heads * self.head_dim,
            hidden_size,
            bias=False,
170
            quant_config=quant_config,
171
            prefix=f"{prefix}.o_proj",
172
173
        )
        # Create the alibi slopes and slice them.
174
        if self.position_embedding == "ALIBI":
175
176
177
178
179
180
181
            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
182
183
184
185
186
187
188
189
            self.attn = Attention(
                self.num_heads,
                self.head_dim,
                scaling,
                alibi_slopes=alibi_slopes,
                quant_config=quant_config,
                prefix=f"{prefix}.attn",
            )
190
191
192
193
194
195
196
197
        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
198
199
200
201
202
203
204
205
            self.attn = Attention(
                self.num_heads,
                self.head_dim,
                self.scaling,
                cache_config=cache_config,
                quant_config=quant_config,
                prefix=f"{prefix}.attn",
            )
206
207
208
209
210
211
212
213

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        qkv, _ = self.W_pack(hidden_states)
        q, k, v = qkv.chunk(chunks=3, dim=-1)
214
        if self.position_embedding != "ALIBI":
215
            q, k = self.rotary_emb(positions, q, k)
216
        attn_output = self.attn(q, k, v)
217
218
219
220
221
        output, _ = self.o_proj(attn_output)
        return output


class BaiChuanDecoderLayer(nn.Module):
222
223
224
225
    def __init__(
        self,
        config: PretrainedConfig,
        position_embedding: str,
226
227
        cache_config: CacheConfig | None = None,
        quant_config: QuantizationConfig | None = None,
228
229
        prefix: str = "",
    ):
230
231
232
        super().__init__()
        self.hidden_size = config.hidden_size
        rope_theta = getattr(config, "rope_theta", 10000)
233
        max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
234
235
236
237
238
239
        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,
240
            cache_config=cache_config,
241
            quant_config=quant_config,
242
            prefix=f"{prefix}.self_attn",
243
244
245
246
247
        )
        self.mlp = BaiChuanMLP(
            hidden_size=self.hidden_size,
            intermediate_size=config.intermediate_size,
            hidden_act=config.hidden_act,
248
            quant_config=quant_config,
249
            prefix=f"{prefix}.mlp",
250
        )
251
252
253
254
        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
        )
255
256
257
258
259

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
260
        residual: torch.Tensor | None,
261
    ) -> tuple[torch.Tensor, torch.Tensor]:
262
263
264
265
266
        # Self Attention
        if residual is None:
            residual = hidden_states
            hidden_states = self.input_layernorm(hidden_states)
        else:
267
            hidden_states, residual = self.input_layernorm(hidden_states, residual)
268
269
270
271
272
273
        hidden_states = self.self_attn(
            positions=positions,
            hidden_states=hidden_states,
        )

        # Fully Connected
274
        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
275
276
277
278
        hidden_states = self.mlp(hidden_states)
        return hidden_states, residual


279
@support_torch_compile
280
class BaiChuanModel(nn.Module):
281
282
283
284
285
286
    def __init__(
        self,
        vllm_config: VllmConfig,
        prefix: str = "",
        position_embedding: str = "ROPE",
    ) -> None:
287
        super().__init__()
288
289
290
291
292

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

293
294
295
296
297
298
299
        self.config = config
        self.vocab_size = config.vocab_size

        self.embed_tokens = VocabParallelEmbedding(
            config.vocab_size,
            config.hidden_size,
        )
300
301
        self.start_layer, self.end_layer, self.layers = make_layers(
            config.num_hidden_layers,
302
303
304
            lambda prefix: BaiChuanDecoderLayer(
                config, position_embedding, cache_config, quant_config, prefix=prefix
            ),
305
306
            prefix=f"{prefix}.layers",
        )
307
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
308
309
310
        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
            ["hidden_states", "residual"], config.hidden_size
        )
311

312
    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
313
314
        return self.embed_tokens(input_ids)

315
316
317
318
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
319
320
321
        intermediate_tensors: IntermediateTensors | None,
        inputs_embeds: torch.Tensor | None = None,
    ) -> torch.Tensor | IntermediateTensors:
322
        if get_pp_group().is_first_rank:
323
324
325
            if inputs_embeds is not None:
                hidden_states = inputs_embeds
            else:
326
                hidden_states = self.embed_input_ids(input_ids)
327
328
329
330
331
            residual = None
        else:
            assert intermediate_tensors is not None
            hidden_states = intermediate_tensors["hidden_states"]
            residual = intermediate_tensors["residual"]
332
        for layer in islice(self.layers, self.start_layer, self.end_layer):
333
334
335
336
337
            hidden_states, residual = layer(
                positions,
                hidden_states,
                residual,
            )
338
        if not get_pp_group().is_last_rank:
339
340
341
342
343
344
            return IntermediateTensors(
                {
                    "hidden_states": hidden_states,
                    "residual": residual,
                }
            )
345
346
347
        hidden_states, _ = self.norm(hidden_states, residual)
        return hidden_states

348
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
349
350
351
352
353
354
        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())
355
        loaded_params: set[str] = set()
356
357
358
359
        for name, loaded_weight in weights:
            if "rotary_emb.inv_freq" in name:
                continue

360
            for param_name, weight_name, shard_id in stacked_params_mapping:
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
                if weight_name not in name:
                    continue
                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
                if is_pp_missing_parameter(name, self):
                    continue
                param = params_dict[name]
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                break
            else:
                # Skip loading extra bias for GPTQ models.
                if name.endswith(".bias") and name not in params_dict:
                    continue
                if is_pp_missing_parameter(name, self):
                    continue
                param = params_dict[name]
380
                weight_loader = getattr(param, "weight_loader", default_weight_loader)
381
382
383
384
                weight_loader(param, loaded_weight)
            loaded_params.add(name)
        return loaded_params

385

386
class BaiChuanBaseForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsQuant):
387
388
389
390
391
392
393
    packed_modules_mapping = {
        "W_pack": ["W_pack"],
        "gate_up_proj": [
            "gate_proj",
            "up_proj",
        ],
    }
394

395
396
    def __init__(
        self,
397
        *,
398
399
400
        vllm_config: VllmConfig,
        prefix: str = "",
        position_embedding: str = "ROPE",
401
    ):
402
        super().__init__()
403
404
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
405

406
        self.config = config
407

408
        self.tp_size = get_tensor_model_parallel_world_size()
409
        self.quant_config = quant_config
410
411
412
413
414
415
416
417
418
419
420
        self.model = BaiChuanModel(
            vllm_config=vllm_config,
            prefix=prefix,
            position_embedding=position_embedding,
        )
        self.lm_head = ParallelLMHead(
            config.vocab_size,
            config.hidden_size,
            quant_config=quant_config,
            prefix=maybe_prefix(prefix, "lm_head"),
        )
421
        self.lm_head.weight.weight_loader = self.lm_head_weight_loader
422
423
        if self.config.tie_word_embeddings:
            self.lm_head.weight = self.model.embed_tokens.weight
424
        self.logits_processor = LogitsProcessor(config.vocab_size)
425
        self.make_empty_intermediate_tensors = (
426
427
            self.model.make_empty_intermediate_tensors
        )
428

429
430
    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.model.embed_input_ids(input_ids)
431

432
433
434
435
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
436
437
438
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
    ) -> torch.Tensor | IntermediateTensors:
439
440
441
        hidden_states = self.model(
            input_ids, positions, intermediate_tensors, inputs_embeds
        )
442
443
        return hidden_states

444
445
446
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
447
    ) -> torch.Tensor | None:
448
        logits = self.logits_processor(self.lm_head, hidden_states)
449
450
        return logits

451
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
452
453
454
        loader = AutoWeightsLoader(self)
        return loader.load_weights(weights)

455
    def lm_head_weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor):
456
457
458
459
460
461
462
463
464
        # Unlike Baichuan, Baichuan2 normalizes the head weights.
        # Refer to:
        # 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)
465
466
467
468
        if self.tp_size > 1:
            row_parallel_weight_loader(param, loaded_weight)
        else:
            default_weight_loader(param, loaded_weight)
469
470


471
class BaichuanForCausalLM(BaiChuanBaseForCausalLM):
472
473
474
    """Baichuan 13B and Baichuan2 7B/13B.
    NOTE: the class name has a lower case 'c'.
    """
475

476
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
477
        config = vllm_config.model_config.hf_config
478
        if config.hidden_size == 4096:  # baichuan2 7b
479
480
481
            super().__init__(
                vllm_config=vllm_config, prefix=prefix, position_embedding="ROPE"
            )
482
        else:  # baichuan 13b, baichuan2 13b
483
484
485
            super().__init__(
                vllm_config=vllm_config, prefix=prefix, position_embedding="ALIBI"
            )
486
487


488
class BaiChuanForCausalLM(BaiChuanBaseForCausalLM):
489
490
491
    """Baichuan 7B.
    NOTE: the class name has an upper case 'C'.
    """
492

493
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
494
495
496
        super().__init__(
            vllm_config=vllm_config, prefix=prefix, position_embedding="ROPE"
        )