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

4
5
# 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
6
# Copyright 2023 The vLLM team.
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
# 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
25
"""Inference-only LLaMA model compatible with HuggingFace weights."""
26

27
from collections.abc import Iterable
28
from itertools import islice
Woosuk Kwon's avatar
Woosuk Kwon committed
29
30
31
32
33

import torch
from torch import nn
from transformers import LlamaConfig

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

62
63
from .adapters import as_embedding_model, as_seq_cls_model
from .interfaces import (
64
    EagleModelMixin,
65
66
67
68
69
    SupportsEagle,
    SupportsEagle3,
    SupportsLoRA,
    SupportsPP,
)
70
71
72
73
74
75
76
77
78
from .utils import (
    AutoWeightsLoader,
    PPMissingLayer,
    extract_layer_index,
    is_pp_missing_parameter,
    make_empty_intermediate_tensors_factory,
    make_layers,
    maybe_prefix,
)
79

Woosuk Kwon's avatar
Woosuk Kwon committed
80
81
82
83

class LlamaMLP(nn.Module):
    def __init__(
        self,
84
85
86
        hidden_size: int,
        intermediate_size: int,
        hidden_act: str,
87
        quant_config: QuantizationConfig | None = None,
88
        bias: bool = False,
89
        prefix: str = "",
90
        reduce_results: bool = True,
91
        disable_tp: bool = False,
92
    ) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
93
        super().__init__()
94
        self.gate_up_proj = MergedColumnParallelLinear(
95
96
            input_size=hidden_size,
            output_sizes=[intermediate_size] * 2,
97
            bias=bias,
98
            quant_config=quant_config,
99
            disable_tp=disable_tp,
100
101
102
103
104
105
106
            prefix=f"{prefix}.gate_up_proj",
        )
        self.down_proj = RowParallelLinear(
            input_size=intermediate_size,
            output_size=hidden_size,
            bias=bias,
            quant_config=quant_config,
107
            reduce_results=reduce_results,
108
            disable_tp=disable_tp,
109
110
            prefix=f"{prefix}.down_proj",
        )
111
        if hidden_act != "silu":
112
113
114
            raise ValueError(
                f"Unsupported activation: {hidden_act}. Only silu is supported for now."
            )
Woosuk Kwon's avatar
Woosuk Kwon committed
115
        self.act_fn = SiluAndMul()
Woosuk Kwon's avatar
Woosuk Kwon committed
116
117

    def forward(self, x):
118
119
        x, _ = self.gate_up_proj(x)
        x = self.act_fn(x)
Woosuk Kwon's avatar
Woosuk Kwon committed
120
121
122
123
124
        x, _ = self.down_proj(x)
        return x


class LlamaAttention(nn.Module):
125
126
127
128
129
130
131
    def __init__(
        self,
        config: LlamaConfig,
        hidden_size: int,
        num_heads: int,
        num_kv_heads: int,
        max_position_embeddings: int = 8192,
132
        quant_config: QuantizationConfig | None = None,
133
134
        bias: bool = False,
        bias_o_proj: bool = False,
135
        cache_config: CacheConfig | None = None,
136
137
138
        prefix: str = "",
        attn_type: str = AttentionType.DECODER,
    ) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
139
        super().__init__()
140
        layer_idx = extract_layer_index(prefix)
141
        self.hidden_size = hidden_size
Zhuohan Li's avatar
Zhuohan Li committed
142
        tp_size = get_tensor_model_parallel_world_size()
143
        self.total_num_heads = num_heads
Zhuohan Li's avatar
Zhuohan Li committed
144
145
        assert self.total_num_heads % tp_size == 0
        self.num_heads = self.total_num_heads // tp_size
146
        self.total_num_kv_heads = num_kv_heads
147
148
149
150
151
152
153
154
155
        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)
156

157
        head_dim = getattr(config, "head_dim", None)
158
        self.head_dim = head_dim or self.hidden_size // self.total_num_heads
Zhuohan Li's avatar
Zhuohan Li committed
159
160
        self.q_size = self.num_heads * self.head_dim
        self.kv_size = self.num_kv_heads * self.head_dim
161
        self.scaling = self.head_dim**-0.5
162
        self.max_position_embeddings = max_position_embeddings
Woosuk Kwon's avatar
Woosuk Kwon committed
163

164
        self.qkv_proj = QKVParallelLinear(
165
166
167
168
            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,
169
            bias=bias,
170
            quant_config=quant_config,
171
            prefix=f"{prefix}.qkv_proj",
Woosuk Kwon's avatar
Woosuk Kwon committed
172
        )
173

174
        self.o_proj = RowParallelLinear(
175
176
            input_size=self.total_num_heads * self.head_dim,
            output_size=hidden_size,
177
            bias=bias_o_proj,
178
            quant_config=quant_config,
179
            prefix=f"{prefix}.o_proj",
Woosuk Kwon's avatar
Woosuk Kwon committed
180
        )
Woosuk Kwon's avatar
Woosuk Kwon committed
181

182
        self._init_rotary_emb(config, quant_config=quant_config)
183

184
185
        sliding_window = None
        if layer_types := getattr(config, "layer_types", None):
186
187
188
            # Fix for Eagle3 compatibility:
            # for draft models, subtract target layer count
            # to get draft-relative layer index starting from 0
189
            if hasattr(config, "target_layer_count"):
190
191
192
193
194
195
                # This is a draft model,
                # adjust layer_idx to be relative to draft layers
                effective_layer_idx = layer_idx - config.target_layer_count
            else:
                # This is a target model, use layer_idx directly
                effective_layer_idx = layer_idx
196
            assert effective_layer_idx < len(layer_types), (
197
198
                f"effective_layer_idx: {effective_layer_idx} "
                f"is out of bounds for layer_types: {layer_types}"
199
            )
200

201
            is_sliding = layer_types[effective_layer_idx] == "sliding_attention"
202
203
            if is_sliding:
                sliding_window = config.sliding_window
204

205
206
207
208
209
        attn_cls = (
            EncoderOnlyAttention
            if attn_type == AttentionType.ENCODER_ONLY
            else Attention
        )
210
211

        self.attn = attn_cls(
212
213
214
215
216
217
            self.num_heads,
            self.head_dim,
            self.scaling,
            num_kv_heads=self.num_kv_heads,
            cache_config=cache_config,
            quant_config=quant_config,
218
            per_layer_sliding_window=sliding_window,
219
            attn_type=attn_type,
220
            prefix=f"{prefix}.attn",
221
        )
Woosuk Kwon's avatar
Woosuk Kwon committed
222
223
224

    def forward(
        self,
225
        positions: torch.Tensor,
Woosuk Kwon's avatar
Woosuk Kwon committed
226
227
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
228
        qkv, _ = self.qkv_proj(hidden_states)
Zhuohan Li's avatar
Zhuohan Li committed
229
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
230
        q, k = self.rotary_emb(positions, q, k)
231
        attn_output = self.attn(q, k, v)
Woosuk Kwon's avatar
Woosuk Kwon committed
232
233
234
        output, _ = self.o_proj(attn_output)
        return output

235
236
237
    def _init_rotary_emb(
        self,
        config: LlamaConfig,
238
        quant_config: QuantizationConfig | None,
239
    ) -> None:
240
241
        is_neox_style = True
        is_gguf = quant_config and quant_config.get_name() == "gguf"
242
        if is_gguf and config.model_type == "llama":
243
244
245
246
247
            is_neox_style = False

        self.rotary_emb = get_rope(
            self.head_dim,
            max_position=self.max_position_embeddings,
248
            rope_parameters=getattr(config, "rope_parameters", None),
249
250
251
            is_neox_style=is_neox_style,
        )

Woosuk Kwon's avatar
Woosuk Kwon committed
252
253

class LlamaDecoderLayer(nn.Module):
254
255
256
257
    def __init__(
        self,
        vllm_config: VllmConfig,
        prefix: str = "",
258
        config: LlamaConfig | None = None,
259
        attn_layer_type: type[nn.Module] = LlamaAttention,
260
    ) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
261
        super().__init__()
262
263
264

        config = config or vllm_config.model_config.hf_config
        cache_config = vllm_config.cache_config
265
        quant_config = self.get_quant_config(vllm_config)
266

Woosuk Kwon's avatar
Woosuk Kwon committed
267
        self.hidden_size = config.hidden_size
268
        max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
269
270
271
        # Support abacusai/Smaug-72B-v0.1 with attention_bias
        # Support internlm/internlm-7b with bias
        attention_bias = getattr(config, "attention_bias", False) or getattr(
272
273
            config, "bias", False
        )
274
275
        bias_o_proj = attention_bias
        # support internlm/internlm3-8b with qkv_bias
276
        if hasattr(config, "qkv_bias"):
277
278
            attention_bias = config.qkv_bias

279
280
281
282
283
284
285
286
287
        # By default, Llama uses causal attention as it is a decoder-only model.
        # You can override the HF config with `is_causal=False` to enable
        # bidirectional attention, which is used in some embedding models
        # (e.g. parasail-ai/GritLM-7B-vllm)
        if getattr(config, "is_causal", True):
            attn_type = AttentionType.DECODER
        else:
            attn_type = AttentionType.ENCODER_ONLY

288
        self.self_attn = attn_layer_type(
289
            config=config,
290
291
            hidden_size=self.hidden_size,
            num_heads=config.num_attention_heads,
292
293
294
            num_kv_heads=getattr(
                config, "num_key_value_heads", config.num_attention_heads
            ),
295
            max_position_embeddings=max_position_embeddings,
296
            quant_config=quant_config,
297
            bias=attention_bias,
298
            bias_o_proj=bias_o_proj,
299
            cache_config=cache_config,
300
            prefix=f"{prefix}.self_attn",
301
            attn_type=attn_type,
Woosuk Kwon's avatar
Woosuk Kwon committed
302
303
        )
        self.mlp = LlamaMLP(
304
305
306
            hidden_size=self.hidden_size,
            intermediate_size=config.intermediate_size,
            hidden_act=config.hidden_act,
307
            quant_config=quant_config,
308
            bias=getattr(config, "mlp_bias", False),
309
            prefix=f"{prefix}.mlp",
Woosuk Kwon's avatar
Woosuk Kwon committed
310
        )
311
312
313
314
        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
315
316
317

    def forward(
        self,
318
        positions: torch.Tensor,
Woosuk Kwon's avatar
Woosuk Kwon committed
319
        hidden_states: torch.Tensor,
320
        residual: torch.Tensor | None,
321
    ) -> tuple[torch.Tensor, torch.Tensor]:
Woosuk Kwon's avatar
Woosuk Kwon committed
322
        # Self Attention
323
324
325
326
        if residual is None:
            residual = hidden_states
            hidden_states = self.input_layernorm(hidden_states)
        else:
327
328
            hidden_states, residual = self.input_layernorm(hidden_states, residual)
        hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states)
Woosuk Kwon's avatar
Woosuk Kwon committed
329
330

        # Fully Connected
331
        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
Woosuk Kwon's avatar
Woosuk Kwon committed
332
        hidden_states = self.mlp(hidden_states)
333
        return hidden_states, residual
Woosuk Kwon's avatar
Woosuk Kwon committed
334

335
    def get_quant_config(self, vllm_config: VllmConfig) -> QuantizationConfig | None:
336
337
338
        """Get quantization config for this layer. Override in subclasses."""
        return vllm_config.quant_config

Woosuk Kwon's avatar
Woosuk Kwon committed
339

340
341
342
343
344
345
346
347
348
349
def llama_model_invariants(
    input_ids, positions, intermediate_tensors=None, inputs_embeds=None
):
    """Shape invariants for Llama model compilation, those are translated to
    runtime assertions for unbacked dynamic shapes and are compiled away for
    backed"""
    if input_ids is not None:
        torch._check(positions.size()[0] == input_ids.size()[0])


350
351
352
353
354
@support_torch_compile(
    # TODO[#32068]: Investigate recompilation
    # mark_unbacked_dims={"input_ids": 0},
    shape_invariants=llama_model_invariants
)
355
class LlamaModel(nn.Module, EagleModelMixin):
356
357
358
359
360
361
362
    def __init__(
        self,
        *,
        vllm_config: VllmConfig,
        prefix: str = "",
        layer_type: type[nn.Module] = LlamaDecoderLayer,
    ):
Woosuk Kwon's avatar
Woosuk Kwon committed
363
        super().__init__()
364
365
366
367

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

Woosuk Kwon's avatar
Woosuk Kwon committed
368
        self.config = config
369
        self.quant_config = quant_config
370
371
372

        self.vocab_size = config.vocab_size

373
374
375
        if get_pp_group().is_first_rank or (
            config.tie_word_embeddings and get_pp_group().is_last_rank
        ):
376
377
378
            self.embed_tokens = VocabParallelEmbedding(
                self.vocab_size,
                config.hidden_size,
379
                quant_config=quant_config,
380
381
382
            )
        else:
            self.embed_tokens = PPMissingLayer()
383
        self.start_layer, self.end_layer, self.layers = make_layers(
384
            config.num_hidden_layers,
385
            lambda prefix: layer_type(vllm_config=vllm_config, prefix=prefix),
386
387
            prefix=f"{prefix}.layers",
        )
388
389
390
391
        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
392

393
394
395
        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
            ["hidden_states", "residual"], config.hidden_size
        )
396

397
    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
398
399
        return self.embed_tokens(input_ids)

Woosuk Kwon's avatar
Woosuk Kwon committed
400
401
    def forward(
        self,
402
        input_ids: torch.Tensor | None,
403
        positions: torch.Tensor,
404
405
        intermediate_tensors: IntermediateTensors | None,
        inputs_embeds: torch.Tensor | None = None,
406
        **extra_layer_kwargs,
407
    ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
408
409
410
411
        if get_pp_group().is_first_rank:
            if inputs_embeds is not None:
                hidden_states = inputs_embeds
            else:
412
                hidden_states = self.embed_input_ids(input_ids)
413
            residual = None
414
        else:
415
416
417
418
            assert intermediate_tensors is not None
            hidden_states = intermediate_tensors["hidden_states"]
            residual = intermediate_tensors["residual"]

419
        aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)
420
        for idx, layer in enumerate(
421
422
            islice(self.layers, self.start_layer, self.end_layer)
        ):
423
424
425
            hidden_states, residual = layer(
                positions, hidden_states, residual, **extra_layer_kwargs
            )
426
427
428
            self._maybe_add_hidden_state(
                aux_hidden_states, idx + 1, hidden_states, residual
            )
429
430

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

435
        hidden_states, _ = self.norm(hidden_states, residual)
436
437
438

        if len(aux_hidden_states) > 0:
            return hidden_states, aux_hidden_states
Woosuk Kwon's avatar
Woosuk Kwon committed
439
440
        return hidden_states

441
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
442
443
444
445
446
447
448
449
450
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
            (".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),
        ]
        params_dict = dict(self.named_parameters())
451
        loaded_params: set[str] = set()
452
453
454
        for name, loaded_weight in weights:
            if "rotary_emb.inv_freq" in name:
                continue
455
            if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
456
457
458
                # Models trained using ColossalAI may include these tensors in
                # the checkpoint. Skip them.
                continue
459
460
461
            if self.quant_config is not None and (
                scale_name := self.quant_config.get_cache_scale(name)
            ):
462
                # Loading kv cache quantization scales
463
                param = params_dict[scale_name]
464
465
466
467
                weight_loader = getattr(param, "weight_loader", default_weight_loader)
                loaded_weight = (
                    loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
                )
468
                weight_loader(param, loaded_weight)
469
                loaded_params.add(scale_name)
470
                continue
471
472
            if "scale" in name or "zero_point" in name:
                # Remapping the name of FP8 kv-scale or zero point.
473
474
475
                name = maybe_remap_kv_scale_name(name, params_dict)
                if name is None:
                    continue
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
            for param_name, weight_name, shard_id in stacked_params_mapping:
                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]
500
                weight_loader = getattr(param, "weight_loader", default_weight_loader)
501
                weight_loader(param, loaded_weight)
502
503
            loaded_params.add(name)
        return loaded_params
504

Woosuk Kwon's avatar
Woosuk Kwon committed
505

506
507
508
class LlamaForCausalLM(
    nn.Module, SupportsLoRA, SupportsPP, SupportsEagle, SupportsEagle3
):
Terry's avatar
Terry committed
509
    packed_modules_mapping = {
510
        "qkv_proj": ["q_proj", "k_proj", "v_proj"],
511
        "gate_up_proj": ["gate_proj", "up_proj"],
Terry's avatar
Terry committed
512
513
514
515
516
    }

    # LoRA specific attributes
    embedding_modules = {
        "embed_tokens": "input_embeddings",
517
        "lm_head": "output_embeddings",
Terry's avatar
Terry committed
518
    }
519

520
521
522
523
524
525
526
    def __init__(
        self,
        *,
        vllm_config: VllmConfig,
        prefix: str = "",
        layer_type: type[nn.Module] = LlamaDecoderLayer,
    ):
Woosuk Kwon's avatar
Woosuk Kwon committed
527
        super().__init__()
528
529
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
Woosuk Kwon's avatar
Woosuk Kwon committed
530
        self.config = config
531

532
533
534
535
536
        self.model = self._init_model(
            vllm_config=vllm_config,
            prefix=maybe_prefix(prefix, "model"),
            layer_type=layer_type,
        )
537

538
539
        if get_pp_group().is_last_rank:
            self.lm_head = ParallelLMHead(
540
                config.vocab_size,
541
542
                config.hidden_size,
                quant_config=quant_config,
543
                prefix=maybe_prefix(prefix, "lm_head"),
544
545
            )
            if config.tie_word_embeddings:
546
                self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)
547
548

            logit_scale = getattr(config, "logit_scale", 1.0)
549
            self.logits_processor = LogitsProcessor(
550
                config.vocab_size, scale=logit_scale
551
            )
552
553
        else:
            self.lm_head = PPMissingLayer()
554

555
        self.make_empty_intermediate_tensors = (
556
557
            self.model.make_empty_intermediate_tensors
        )
Woosuk Kwon's avatar
Woosuk Kwon committed
558

559
560
561
562
563
564
565
    def _init_model(
        self,
        vllm_config: VllmConfig,
        prefix: str = "",
        layer_type: type[nn.Module] = LlamaDecoderLayer,
    ):
        return LlamaModel(vllm_config=vllm_config, prefix=prefix, layer_type=layer_type)
566

567
568
    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.model.embed_input_ids(input_ids)
569

Woosuk Kwon's avatar
Woosuk Kwon committed
570
571
    def forward(
        self,
zhuwenwen's avatar
zhuwenwen committed
572
        input_ids: torch.Tensor | None,
573
        positions: torch.Tensor,
574
575
576
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
    ) -> torch.Tensor | IntermediateTensors:
577
578
579
        model_output = self.model(
            input_ids, positions, intermediate_tensors, inputs_embeds
        )
580
        return model_output
581

582
583
584
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
585
    ) -> torch.Tensor | None:
586
        logits = self.logits_processor(self.lm_head, hidden_states)
587
588
        return logits

589
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
590
591
        loader = AutoWeightsLoader(
            self,
592
            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
593
        )
594
        return loader.load_weights(weights)
595
596
597
598
599
600
601
602
603
604
605


class LlamaBidirectionalForSequenceClassification(as_seq_cls_model(LlamaForCausalLM)):
    # This class sets the correct attention type and pooling type
    # through LlamaBidirectionalConfig.
    pass


class LlamaBidirectionalModel(as_embedding_model(LlamaForCausalLM)):
    # This class sets the correct attention type and pooling type
    # through LlamaBidirectionalConfig.
zhuwenwen's avatar
zhuwenwen committed
606
    pass