llama.py 20.7 KB
Newer Older
1
# coding=utf-8
2
3
# Adapted from
# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/llama/modeling_llama.py
Woosuk Kwon's avatar
Woosuk Kwon committed
4
# Copyright 2023 The vLLM team.
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
# 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
23
"""Inference-only LLaMA model compatible with HuggingFace weights."""
24
from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
Woosuk Kwon's avatar
Woosuk Kwon committed
25
26
27
28
29

import torch
from torch import nn
from transformers import LlamaConfig

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

52
from .interfaces import SupportsLoRA
53
from .utils import PPMissingLayer, is_pp_missing_parameter, make_layers
54

Woosuk Kwon's avatar
Woosuk Kwon committed
55
56

class LlamaMLP(nn.Module):
57

Woosuk Kwon's avatar
Woosuk Kwon committed
58
59
    def __init__(
        self,
60
61
62
        hidden_size: int,
        intermediate_size: int,
        hidden_act: str,
63
        quant_config: Optional[QuantizationConfig] = None,
64
        bias: bool = False,
65
        prefix: str = "",
66
    ) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
67
        super().__init__()
68
        self.gate_up_proj = MergedColumnParallelLinear(
69
70
            input_size=hidden_size,
            output_sizes=[intermediate_size] * 2,
71
            bias=bias,
72
73
            quant_config=quant_config,
            prefix=f"{prefix}.gate_up_proj")
74
75
        self.down_proj = RowParallelLinear(input_size=intermediate_size,
                                           output_size=hidden_size,
76
                                           bias=bias,
77
78
                                           quant_config=quant_config,
                                           prefix=f"{prefix}.down_proj")
79
80
81
        if hidden_act != "silu":
            raise ValueError(f"Unsupported activation: {hidden_act}. "
                             "Only silu is supported for now.")
Woosuk Kwon's avatar
Woosuk Kwon committed
82
        self.act_fn = SiluAndMul()
Woosuk Kwon's avatar
Woosuk Kwon committed
83
84

    def forward(self, x):
85
        gate_up, _ = self.gate_up_proj(x)
Woosuk Kwon's avatar
Woosuk Kwon committed
86
        x = self.act_fn(gate_up)
Woosuk Kwon's avatar
Woosuk Kwon committed
87
88
89
90
91
92
93
94
        x, _ = self.down_proj(x)
        return x


class LlamaAttention(nn.Module):

    def __init__(
        self,
95
        config: LlamaConfig,
96
97
98
99
100
101
        hidden_size: int,
        num_heads: int,
        num_kv_heads: int,
        rope_theta: float = 10000,
        rope_scaling: Optional[Dict[str, Any]] = None,
        max_position_embeddings: int = 8192,
102
        quant_config: Optional[QuantizationConfig] = None,
103
        bias: bool = False,
104
        cache_config: Optional[CacheConfig] = None,
105
        prefix: str = "",
106
    ) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
107
        super().__init__()
108
        self.hidden_size = hidden_size
Zhuohan Li's avatar
Zhuohan Li committed
109
        tp_size = get_tensor_model_parallel_world_size()
110
        self.total_num_heads = num_heads
Zhuohan Li's avatar
Zhuohan Li committed
111
112
        assert self.total_num_heads % tp_size == 0
        self.num_heads = self.total_num_heads // tp_size
113
        self.total_num_kv_heads = num_kv_heads
114
115
116
117
118
119
120
121
122
        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)
123
124
125
        # MistralConfig has an optional head_dim introduced by Mistral-Nemo
        self.head_dim = getattr(config, "head_dim",
                                self.hidden_size // self.total_num_heads)
Zhuohan Li's avatar
Zhuohan Li committed
126
127
        self.q_size = self.num_heads * self.head_dim
        self.kv_size = self.num_kv_heads * self.head_dim
128
        self.scaling = self.head_dim**-0.5
129
130
        self.rope_theta = rope_theta
        self.max_position_embeddings = max_position_embeddings
Woosuk Kwon's avatar
Woosuk Kwon committed
131

132
        self.qkv_proj = QKVParallelLinear(
133
134
135
136
            hidden_size=hidden_size,
            head_size=self.head_dim,
            total_num_heads=self.total_num_heads,
            total_num_kv_heads=self.total_num_kv_heads,
137
            bias=bias,
138
            quant_config=quant_config,
139
            prefix=f"{prefix}.qkv_proj",
Woosuk Kwon's avatar
Woosuk Kwon committed
140
        )
141
        self.o_proj = RowParallelLinear(
142
143
            input_size=self.total_num_heads * self.head_dim,
            output_size=hidden_size,
144
            bias=bias,
145
            quant_config=quant_config,
146
            prefix=f"{prefix}.o_proj",
Woosuk Kwon's avatar
Woosuk Kwon committed
147
        )
Woosuk Kwon's avatar
Woosuk Kwon committed
148

149
150
151
152
153
154
155
        self.rotary_emb = get_rope(
            self.head_dim,
            rotary_dim=self.head_dim,
            max_position=max_position_embeddings,
            base=rope_theta,
            rope_scaling=rope_scaling,
        )
156
157
158
159
        self.attn = Attention(self.num_heads,
                              self.head_dim,
                              self.scaling,
                              num_kv_heads=self.num_kv_heads,
160
161
                              cache_config=cache_config,
                              quant_config=quant_config)
Woosuk Kwon's avatar
Woosuk Kwon committed
162
163
164

    def forward(
        self,
165
        positions: torch.Tensor,
Woosuk Kwon's avatar
Woosuk Kwon committed
166
        hidden_states: torch.Tensor,
167
168
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
Woosuk Kwon's avatar
Woosuk Kwon committed
169
    ) -> torch.Tensor:
170
        qkv, _ = self.qkv_proj(hidden_states)
Zhuohan Li's avatar
Zhuohan Li committed
171
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
172
        q, k = self.rotary_emb(positions, q, k)
173
        attn_output = self.attn(q, k, v, kv_cache, attn_metadata)
Woosuk Kwon's avatar
Woosuk Kwon committed
174
175
176
177
178
179
        output, _ = self.o_proj(attn_output)
        return output


class LlamaDecoderLayer(nn.Module):

180
181
182
    def __init__(
        self,
        config: LlamaConfig,
183
        cache_config: Optional[CacheConfig] = None,
184
        quant_config: Optional[QuantizationConfig] = None,
185
        prefix: str = "",
186
    ) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
187
188
        super().__init__()
        self.hidden_size = config.hidden_size
189
190
        rope_theta = getattr(config, "rope_theta", 10000)
        rope_scaling = getattr(config, "rope_scaling", None)
191
192
193
194
        if rope_scaling is not None and getattr(
                config, "original_max_position_embeddings", None):
            rope_scaling["original_max_position_embeddings"] = (
                config.original_max_position_embeddings)
195
196
        max_position_embeddings = getattr(config, "max_position_embeddings",
                                          8192)
197
198
199
200
        # Support abacusai/Smaug-72B-v0.1 with attention_bias
        # Support internlm/internlm-7b with bias
        attention_bias = getattr(config, "attention_bias", False) or getattr(
            config, "bias", False)
Woosuk Kwon's avatar
Woosuk Kwon committed
201
        self.self_attn = LlamaAttention(
202
            config=config,
203
204
            hidden_size=self.hidden_size,
            num_heads=config.num_attention_heads,
205
206
            num_kv_heads=getattr(config, "num_key_value_heads",
                                 config.num_attention_heads),
207
208
209
            rope_theta=rope_theta,
            rope_scaling=rope_scaling,
            max_position_embeddings=max_position_embeddings,
210
            quant_config=quant_config,
211
            bias=attention_bias,
212
            cache_config=cache_config,
213
            prefix=f"{prefix}.self_attn",
Woosuk Kwon's avatar
Woosuk Kwon committed
214
215
        )
        self.mlp = LlamaMLP(
216
217
218
            hidden_size=self.hidden_size,
            intermediate_size=config.intermediate_size,
            hidden_act=config.hidden_act,
219
            quant_config=quant_config,
220
            bias=getattr(config, "mlp_bias", False),
221
            prefix=f"{prefix}.mlp",
Woosuk Kwon's avatar
Woosuk Kwon committed
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)
Woosuk Kwon's avatar
Woosuk Kwon committed
227
228
229

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

        # Fully Connected
251
252
        hidden_states, residual = self.post_attention_layernorm(
            hidden_states, residual)
Woosuk Kwon's avatar
Woosuk Kwon committed
253
        hidden_states = self.mlp(hidden_states)
254
        return hidden_states, residual
Woosuk Kwon's avatar
Woosuk Kwon committed
255
256
257
258


class LlamaModel(nn.Module):

259
260
261
    def __init__(
        self,
        config: LlamaConfig,
262
        cache_config: Optional[CacheConfig] = None,
263
        quant_config: Optional[QuantizationConfig] = None,
264
        lora_config: Optional[LoRAConfig] = None,
265
        prefix: str = "",
266
    ) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
267
268
269
        super().__init__()
        self.config = config
        self.padding_idx = config.pad_token_id
270
271
272
273
        lora_vocab = (lora_config.lora_extra_vocab_size *
                      (lora_config.max_loras or 1)) if lora_config else 0
        self.vocab_size = config.vocab_size + lora_vocab
        self.org_vocab_size = config.vocab_size
274
275
276
277
278
279
280
281
282
        if get_pp_group().is_first_rank or (config.tie_word_embeddings
                                            and get_pp_group().is_last_rank):
            self.embed_tokens = VocabParallelEmbedding(
                self.vocab_size,
                config.hidden_size,
                org_num_embeddings=config.vocab_size,
            )
        else:
            self.embed_tokens = PPMissingLayer()
283
        self.start_layer, self.end_layer, self.layers = make_layers(
284
            config.num_hidden_layers,
285
286
287
288
289
            lambda prefix: LlamaDecoderLayer(config=config,
                                             cache_config=cache_config,
                                             quant_config=quant_config,
                                             prefix=prefix),
            prefix=f"{prefix}.layers")
290
291
292
293
        if get_pp_group().is_last_rank:
            self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        else:
            self.norm = PPMissingLayer()
Woosuk Kwon's avatar
Woosuk Kwon committed
294

295
296
297
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embed_tokens(input_ids)

Woosuk Kwon's avatar
Woosuk Kwon committed
298
299
    def forward(
        self,
300
        input_ids: Optional[torch.Tensor],
301
        positions: torch.Tensor,
302
303
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
304
        intermediate_tensors: Optional[IntermediateTensors],
305
        inputs_embeds: Optional[torch.Tensor] = None,
306
307
308
309
310
311
312
    ) -> Union[torch.Tensor, IntermediateTensors]:
        if get_pp_group().is_first_rank:
            if inputs_embeds is not None:
                hidden_states = inputs_embeds
            else:
                hidden_states = self.get_input_embeddings(input_ids)
            residual = None
313
        else:
314
315
316
317
318
            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):
Woosuk Kwon's avatar
Woosuk Kwon committed
319
            layer = self.layers[i]
320
            hidden_states, residual = layer(
Woosuk Kwon's avatar
Woosuk Kwon committed
321
322
                positions,
                hidden_states,
323
                kv_caches[i - self.start_layer],
324
                attn_metadata,
325
                residual,
Woosuk Kwon's avatar
Woosuk Kwon committed
326
            )
327
328
329
330
331
332
333

        if not get_pp_group().is_last_rank:
            return IntermediateTensors({
                "hidden_states": hidden_states,
                "residual": residual
            })

334
        hidden_states, _ = self.norm(hidden_states, residual)
Woosuk Kwon's avatar
Woosuk Kwon committed
335
336
337
        return hidden_states


338
class LlamaForCausalLM(nn.Module, SupportsLoRA):
Terry's avatar
Terry committed
339
340
341
342
343
344
345
346
347
348
349
350
351
352
    packed_modules_mapping = {
        "qkv_proj": [
            "q_proj",
            "k_proj",
            "v_proj",
        ],
        "gate_up_proj": [
            "gate_proj",
            "up_proj",
        ],
    }

    # LoRA specific attributes
    supported_lora_modules = [
353
354
        "qkv_proj", "o_proj", "gate_up_proj", "down_proj", "embed_tokens",
        "lm_head"
Terry's avatar
Terry committed
355
356
357
358
359
360
    ]
    embedding_modules = {
        "embed_tokens": "input_embeddings",
        "lm_head": "output_embeddings",
    }
    embedding_padding_modules = ["lm_head"]
361
362
363
364
365
366
367
368
    bitsandbytes_stacked_params_mapping = {
        # shard_name, weight_name, index
        "q_proj": ("qkv_proj", 0),
        "k_proj": ("qkv_proj", 1),
        "v_proj": ("qkv_proj", 2),
        "gate_proj": ("gate_up_proj", 0),
        "up_proj": ("gate_up_proj", 1),
    }
369

370
371
372
    def __init__(
        self,
        config: LlamaConfig,
373
        cache_config: Optional[CacheConfig] = None,
374
        quant_config: Optional[QuantizationConfig] = None,
375
        lora_config: Optional[LoRAConfig] = None,
376
    ) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
377
        super().__init__()
378

Woosuk Kwon's avatar
Woosuk Kwon committed
379
        self.config = config
380
381
        self.lora_config = lora_config

382
383
384
        self.model = LlamaModel(config,
                                cache_config,
                                quant_config,
385
386
                                lora_config=lora_config,
                                prefix="model")
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
        if get_pp_group().is_last_rank:
            self.unpadded_vocab_size = config.vocab_size
            if lora_config:
                self.unpadded_vocab_size += lora_config.lora_extra_vocab_size
            self.lm_head = ParallelLMHead(
                self.unpadded_vocab_size,
                config.hidden_size,
                org_num_embeddings=config.vocab_size,
                padding_size=DEFAULT_VOCAB_PADDING_SIZE
                # We need bigger padding if using lora for kernel
                # compatibility
                if not lora_config else lora_config.lora_vocab_padding_size,
                quant_config=quant_config,
            )
            if config.tie_word_embeddings:
                self.lm_head.weight = self.model.embed_tokens.weight

            logit_scale = getattr(config, "logit_scale", 1.0)
            self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
                                                    config.vocab_size,
                                                    logit_scale)
            self.sampler = Sampler()
        else:
            self.lm_head = PPMissingLayer()
Woosuk Kwon's avatar
Woosuk Kwon committed
411
412
413

    def forward(
        self,
414
415
        input_ids: torch.Tensor,
        positions: torch.Tensor,
416
417
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
418
419
420
421
422
        intermediate_tensors: Optional[IntermediateTensors] = None,
    ) -> Union[torch.Tensor, IntermediateTensors]:
        model_output = self.model(input_ids, positions, kv_caches,
                                  attn_metadata, intermediate_tensors)
        return model_output
423

424
425
    def compute_logits(self, hidden_states: torch.Tensor,
                       sampling_metadata: SamplingMetadata) -> torch.Tensor:
426
        logits = self.logits_processor(self.lm_head, hidden_states,
427
428
429
                                       sampling_metadata)
        return logits

430
431
    def sample(
        self,
432
        logits: torch.Tensor,
433
        sampling_metadata: SamplingMetadata,
434
    ) -> Optional[SamplerOutput]:
435
        next_tokens = self.sampler(logits, sampling_metadata)
Woosuk Kwon's avatar
Woosuk Kwon committed
436
437
        return next_tokens

438
439
440
441
442
443
444
445
446
447
448
449
450
451
    def make_empty_intermediate_tensors(
            self, batch_size: int, dtype: torch.dtype,
            device: torch.device) -> IntermediateTensors:
        return IntermediateTensors({
            "hidden_states":
            torch.zeros((batch_size, self.config.hidden_size),
                        dtype=dtype,
                        device=device),
            "residual":
            torch.zeros((batch_size, self.config.hidden_size),
                        dtype=dtype,
                        device=device),
        })

452
    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
453
454
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
455
456
457
458
459
            (".qkv_proj", ".q_proj", "q"),
            (".qkv_proj", ".k_proj", "k"),
            (".qkv_proj", ".v_proj", "v"),
            (".gate_up_proj", ".gate_proj", 0),
            (".gate_up_proj", ".up_proj", 1),
Zhuohan Li's avatar
Zhuohan Li committed
460
        ]
461
        params_dict = dict(self.named_parameters())
462
        for name, loaded_weight in weights:
463
464
            if "rotary_emb.inv_freq" in name:
                continue
465
466
467
468
            if ("rotary_emb.cos_cached" in name
                    or "rotary_emb.sin_cached" in name):
                # Models trained using ColossalAI may include these tensors in
                # the checkpoint. Skip them.
469
                continue
470
            for (param_name, weight_name, shard_id) in stacked_params_mapping:
Zhuohan Li's avatar
Zhuohan Li committed
471
                if weight_name not in name:
472
                    continue
CHU Tianxiang's avatar
CHU Tianxiang committed
473
474
475
476
                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
477
478
479
480
481
482
483
484

                if is_pp_missing_parameter(name, self):
                    continue

                param = params_dict[name]
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)

485
                break
486
            else:
CHU Tianxiang's avatar
CHU Tianxiang committed
487
488
489
                # Skip loading extra bias for GPTQ models.
                if name.endswith(".bias") and name not in params_dict:
                    continue
490
                # Remapping the name of FP8 kv-scale.
491
492
493
                name = maybe_remap_kv_scale_name(name, params_dict)
                if name is None:
                    continue
494
495
496
497
498
499
500
501

                if is_pp_missing_parameter(name, self):
                    continue

                param = params_dict[name]
                weight_loader = getattr(param, "weight_loader",
                                        default_weight_loader)
                weight_loader(param, loaded_weight)
502
503
504
505
506
507
508
509
510
511
512

    # If this function is called, it should always initialize KV cache scale
    # factors (or else raise an exception). Thus, handled exceptions should
    # make sure to leave KV cache scale factors in a known good (dummy) state
    def load_kv_cache_scales(self, quantization_param_path: str) -> None:
        tp_size = get_tensor_model_parallel_world_size()
        tp_rank = get_tensor_model_parallel_rank()
        for layer_idx, scaling_factor in kv_cache_scales_loader(
                quantization_param_path, tp_rank, tp_size,
                self.config.num_hidden_layers,
                self.config.__class__.model_type):
513
514
            if not isinstance(self.model.layers[layer_idx], nn.Identity):
                layer_self_attn = self.model.layers[layer_idx].self_attn
515
516
517
518
519
520
521
522

            if is_hip():
                # The scaling factor convention we are assuming is
                # quantized_value * scaling_factor ~= true_value
                # which is consistent with the practice of setting
                # scaling_factor = tensor_amax / FPtype_max
                scaling_factor *= 2
            if hasattr(layer_self_attn, "kv_scale"):
523
                layer_self_attn.attn._kv_scale = scaling_factor
524
525
526
            else:
                raise RuntimeError("Self attention has no KV cache scaling "
                                   "factor attribute!")