llama.py 18.3 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
Woosuk Kwon's avatar
Woosuk Kwon committed
25
26
27
28

import torch
from torch import nn
from transformers import LlamaConfig
zhuwenwen's avatar
zhuwenwen committed
29
import os
Woosuk Kwon's avatar
Woosuk Kwon committed
30

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


class LlamaMLP(nn.Module):
55

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

    def forward(self, x):
80
        gate_up, _ = self.gate_up_proj(x)
Woosuk Kwon's avatar
Woosuk Kwon committed
81
        x = self.act_fn(gate_up)
Woosuk Kwon's avatar
Woosuk Kwon committed
82
83
84
85
86
87
88
89
        x, _ = self.down_proj(x)
        return x


class LlamaAttention(nn.Module):

    def __init__(
        self,
90
91
92
93
94
95
        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,
96
        quant_config: Optional[QuantizationConfig] = None,
97
        bias: bool = False,
98
        cache_config: Optional[CacheConfig] = None,
99
    ) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
100
        super().__init__()
101
        self.hidden_size = hidden_size
Zhuohan Li's avatar
Zhuohan Li committed
102
        tp_size = get_tensor_model_parallel_world_size()
103
        self.total_num_heads = num_heads
Zhuohan Li's avatar
Zhuohan Li committed
104
105
        assert self.total_num_heads % tp_size == 0
        self.num_heads = self.total_num_heads // tp_size
106
        self.total_num_kv_heads = num_kv_heads
107
108
109
110
111
112
113
114
115
        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)
116
        self.head_dim = hidden_size // self.total_num_heads
Zhuohan Li's avatar
Zhuohan Li committed
117
118
        self.q_size = self.num_heads * self.head_dim
        self.kv_size = self.num_kv_heads * self.head_dim
119
        self.scaling = self.head_dim**-0.5
120
121
        self.rope_theta = rope_theta
        self.max_position_embeddings = max_position_embeddings
Woosuk Kwon's avatar
Woosuk Kwon committed
122

123
        self.qkv_proj = QKVParallelLinear(
124
125
126
127
            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,
128
            bias=bias,
129
            quant_config=quant_config,
Woosuk Kwon's avatar
Woosuk Kwon committed
130
        )
131
        self.o_proj = RowParallelLinear(
132
133
            input_size=self.total_num_heads * self.head_dim,
            output_size=hidden_size,
134
            bias=bias,
135
            quant_config=quant_config,
Woosuk Kwon's avatar
Woosuk Kwon committed
136
        )
Woosuk Kwon's avatar
Woosuk Kwon committed
137

138
139
140
141
142
143
144
        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,
        )
145
146
147
148
        self.attn = Attention(self.num_heads,
                              self.head_dim,
                              self.scaling,
                              num_kv_heads=self.num_kv_heads,
149
150
                              cache_config=cache_config,
                              quant_config=quant_config)
Woosuk Kwon's avatar
Woosuk Kwon committed
151
152
153

    def forward(
        self,
154
        positions: torch.Tensor,
Woosuk Kwon's avatar
Woosuk Kwon committed
155
        hidden_states: torch.Tensor,
156
157
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
Woosuk Kwon's avatar
Woosuk Kwon committed
158
    ) -> torch.Tensor:
159
        qkv, _ = self.qkv_proj(hidden_states)
zhuwenwen's avatar
zhuwenwen committed
160
161
        if os.environ.get('FA_PAD') == '1':
            qkv = qkv[...,:-32]
Zhuohan Li's avatar
Zhuohan Li committed
162
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
163
        q, k = self.rotary_emb(positions, q, k)
164
        attn_output = self.attn(q, k, v, kv_cache, attn_metadata)
Woosuk Kwon's avatar
Woosuk Kwon committed
165
166
167
168
169
170
        output, _ = self.o_proj(attn_output)
        return output


class LlamaDecoderLayer(nn.Module):

171
172
173
    def __init__(
        self,
        config: LlamaConfig,
174
        cache_config: Optional[CacheConfig] = None,
175
        quant_config: Optional[QuantizationConfig] = None,
176
    ) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
177
178
        super().__init__()
        self.hidden_size = config.hidden_size
179
180
        rope_theta = getattr(config, "rope_theta", 10000)
        rope_scaling = getattr(config, "rope_scaling", None)
181
182
183
184
        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)
185
186
        max_position_embeddings = getattr(config, "max_position_embeddings",
                                          8192)
187
188
189
190
        # 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
191
        self.self_attn = LlamaAttention(
192
193
            hidden_size=self.hidden_size,
            num_heads=config.num_attention_heads,
194
195
            num_kv_heads=getattr(config, "num_key_value_heads",
                                 config.num_attention_heads),
196
197
198
            rope_theta=rope_theta,
            rope_scaling=rope_scaling,
            max_position_embeddings=max_position_embeddings,
199
            quant_config=quant_config,
200
            bias=attention_bias,
201
            cache_config=cache_config,
Woosuk Kwon's avatar
Woosuk Kwon committed
202
203
        )
        self.mlp = LlamaMLP(
204
205
206
            hidden_size=self.hidden_size,
            intermediate_size=config.intermediate_size,
            hidden_act=config.hidden_act,
207
            quant_config=quant_config,
208
            bias=getattr(config, "mlp_bias", False),
Woosuk Kwon's avatar
Woosuk Kwon committed
209
        )
210
211
212
213
        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
214
215
216

    def forward(
        self,
217
        positions: torch.Tensor,
Woosuk Kwon's avatar
Woosuk Kwon committed
218
        hidden_states: torch.Tensor,
219
220
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
221
222
        residual: Optional[torch.Tensor],
    ) -> Tuple[torch.Tensor, torch.Tensor]:
Woosuk Kwon's avatar
Woosuk Kwon committed
223
        # Self Attention
224
225
226
227
228
229
        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
230
231
232
233
        hidden_states = self.self_attn(
            positions=positions,
            hidden_states=hidden_states,
            kv_cache=kv_cache,
234
            attn_metadata=attn_metadata,
Woosuk Kwon's avatar
Woosuk Kwon committed
235
236
237
        )

        # Fully Connected
238
239
        hidden_states, residual = self.post_attention_layernorm(
            hidden_states, residual)
Woosuk Kwon's avatar
Woosuk Kwon committed
240
        hidden_states = self.mlp(hidden_states)
241
        return hidden_states, residual
Woosuk Kwon's avatar
Woosuk Kwon committed
242
243
244
245


class LlamaModel(nn.Module):

246
247
248
    def __init__(
        self,
        config: LlamaConfig,
249
        cache_config: Optional[CacheConfig] = None,
250
        quant_config: Optional[QuantizationConfig] = None,
251
        lora_config: Optional[LoRAConfig] = None,
252
    ) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
253
254
255
        super().__init__()
        self.config = config
        self.padding_idx = config.pad_token_id
256
257
258
259
        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
260
        self.embed_tokens = VocabParallelEmbedding(
261
            self.vocab_size,
262
            config.hidden_size,
263
            org_num_embeddings=config.vocab_size,
264
        )
265
        self.layers = nn.ModuleList([
266
267
268
269
            LlamaDecoderLayer(config=config,
                              cache_config=cache_config,
                              quant_config=quant_config)
            for idx in range(config.num_hidden_layers)
270
        ])
271
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
Woosuk Kwon's avatar
Woosuk Kwon committed
272

273
274
275
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embed_tokens(input_ids)

Woosuk Kwon's avatar
Woosuk Kwon committed
276
277
    def forward(
        self,
278
        input_ids: Optional[torch.Tensor],
279
        positions: torch.Tensor,
280
281
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
282
        inputs_embeds: Optional[torch.Tensor] = None,
Woosuk Kwon's avatar
Woosuk Kwon committed
283
    ) -> torch.Tensor:
284
285
286
287
        if inputs_embeds is not None:
            hidden_states = inputs_embeds
        else:
            hidden_states = self.get_input_embeddings(input_ids)
288
        residual = None
Woosuk Kwon's avatar
Woosuk Kwon committed
289
290
        for i in range(len(self.layers)):
            layer = self.layers[i]
291
            hidden_states, residual = layer(
Woosuk Kwon's avatar
Woosuk Kwon committed
292
293
294
                positions,
                hidden_states,
                kv_caches[i],
295
                attn_metadata,
296
                residual,
Woosuk Kwon's avatar
Woosuk Kwon committed
297
            )
298
        hidden_states, _ = self.norm(hidden_states, residual)
Woosuk Kwon's avatar
Woosuk Kwon committed
299
300
301
302
        return hidden_states


class LlamaForCausalLM(nn.Module):
Terry's avatar
Terry committed
303
304
305
306
307
308
309
310
311
312
313
314
315
316
    packed_modules_mapping = {
        "qkv_proj": [
            "q_proj",
            "k_proj",
            "v_proj",
        ],
        "gate_up_proj": [
            "gate_proj",
            "up_proj",
        ],
    }

    # LoRA specific attributes
    supported_lora_modules = [
317
318
        "qkv_proj", "o_proj", "gate_up_proj", "down_proj", "embed_tokens",
        "lm_head"
Terry's avatar
Terry committed
319
320
321
322
323
324
    ]
    embedding_modules = {
        "embed_tokens": "input_embeddings",
        "lm_head": "output_embeddings",
    }
    embedding_padding_modules = ["lm_head"]
325
326
327
328
329
330
331
332
    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),
    }
333

334
335
336
    def __init__(
        self,
        config: LlamaConfig,
337
        cache_config: Optional[CacheConfig] = None,
338
        quant_config: Optional[QuantizationConfig] = None,
339
        lora_config: Optional[LoRAConfig] = None,
340
    ) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
341
342
        super().__init__()
        self.config = config
343
344
345
346
        self.model = LlamaModel(config,
                                cache_config,
                                quant_config,
                                lora_config=lora_config)
Terry's avatar
Terry committed
347
        self.unpadded_vocab_size = config.vocab_size
348
        if lora_config:
Terry's avatar
Terry committed
349
            self.unpadded_vocab_size += lora_config.lora_extra_vocab_size
350
        self.lm_head = ParallelLMHead(
Terry's avatar
Terry committed
351
            self.unpadded_vocab_size,
352
353
354
355
356
357
358
            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,
        )
359
360
        if config.tie_word_embeddings:
            self.lm_head.weight = self.model.embed_tokens.weight
361
362
363
364
365

        logit_scale = getattr(config, "logit_scale", 1.0)
        self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
                                                config.vocab_size, logit_scale)
        self.sampler = Sampler()
Woosuk Kwon's avatar
Woosuk Kwon committed
366
367
368

    def forward(
        self,
369
370
        input_ids: torch.Tensor,
        positions: torch.Tensor,
371
372
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
373
    ) -> torch.Tensor:
374
        hidden_states = self.model(input_ids, positions, kv_caches,
375
                                   attn_metadata)
376
377
        return hidden_states

378
379
380
381
382
383
    def compute_logits(self, hidden_states: torch.Tensor,
                       sampling_metadata: SamplingMetadata) -> torch.Tensor:
        logits = self.logits_processor(self.lm_head.weight, hidden_states,
                                       sampling_metadata)
        return logits

384
385
    def sample(
        self,
386
        logits: torch.Tensor,
387
        sampling_metadata: SamplingMetadata,
388
    ) -> Optional[SamplerOutput]:
389
        next_tokens = self.sampler(logits, sampling_metadata)
Woosuk Kwon's avatar
Woosuk Kwon committed
390
391
        return next_tokens

392
    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
393
394
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
395
396
397
398
399
            (".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
400
        ]
401
        params_dict = dict(self.named_parameters())
402
        for name, loaded_weight in weights:
403
404
            if "rotary_emb.inv_freq" in name:
                continue
405
406
407
408
            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.
409
                continue
410
            for (param_name, weight_name, shard_id) in stacked_params_mapping:
Zhuohan Li's avatar
Zhuohan Li committed
411
                if weight_name not in name:
412
                    continue
CHU Tianxiang's avatar
CHU Tianxiang committed
413
414
415
416
417
                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]
418
419
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
420
                break
421
            else:
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
426
427
428
429
430
431
432
433
434
435
436
437
                # Remapping the name of FP8 kv-scale.
                if name.endswith("kv_scale"):
                    remapped_kv_scale_name = name.replace(
                        ".kv_scale", ".attn.kv_scale")
                    if remapped_kv_scale_name not in params_dict:
                        print_warning_once(
                            f"Found kv scale in the checkpoint (e.g. {name}), "
                            "but not found the expected name in the model "
                            f"(e.g. {remapped_kv_scale_name}). kv-scale is "
                            "not loaded.")
                        continue
                    else:
                        name = remapped_kv_scale_name
438
439
440
441
                param = params_dict[name]
                weight_loader = getattr(param, "weight_loader",
                                        default_weight_loader)
                weight_loader(param, loaded_weight)
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461

    # 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):
            layer_self_attn = self.model.layers[layer_idx].self_attn

            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"):
462
                layer_self_attn.attn._kv_scale = scaling_factor
463
464
465
            else:
                raise RuntimeError("Self attention has no KV cache scaling "
                                   "factor attribute!")