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

张大成's avatar
张大成 committed
4
5
6
7
8
# Adapted from
# https://huggingface.co/OrionStarAI/Orion-14B-Base/blob/main/modeling_orion.py
# Copyright (c) OrionStar Inc.
# LICENSE: https://huggingface.co/OrionStarAI/Orion-14B-Base/blob/main/LICENSE
"""Inference-only Orion-14B model compatible with HuggingFace weights."""
9
from collections.abc import Iterable
10
from itertools import islice
11
from typing import Any, Optional, Union
张大成's avatar
张大成 committed
12
13
14
15
16

import torch
from torch import nn
from transformers import PretrainedConfig

17
from vllm.attention import Attention
18
from vllm.compilation.decorators import support_torch_compile
19
from vllm.config import CacheConfig, VllmConfig
20
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
张大成's avatar
张大成 committed
21
from vllm.model_executor.layers.activation import SiluAndMul
22
from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
张大成's avatar
张大成 committed
23
24
                                               QKVParallelLinear,
                                               RowParallelLinear)
25
from vllm.model_executor.layers.logits_processor import LogitsProcessor
26
from vllm.model_executor.layers.quantization import QuantizationConfig
27
from vllm.model_executor.layers.rotary_embedding import get_rope
张大成's avatar
张大成 committed
28
from vllm.model_executor.layers.vocab_parallel_embedding import (
29
    ParallelLMHead, VocabParallelEmbedding)
30
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
张大成's avatar
张大成 committed
31
from vllm.model_executor.sampling_metadata import SamplingMetadata
32
from vllm.sequence import IntermediateTensors
张大成's avatar
张大成 committed
33

34
from .interfaces import SupportsPP
35
from .utils import (AutoWeightsLoader, is_pp_missing_parameter,
36
37
                    make_empty_intermediate_tensors_factory, make_layers,
                    maybe_prefix)
38

张大成's avatar
张大成 committed
39
40
41
42
43
44
45
46

class OrionMLP(nn.Module):

    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
        hidden_act: str,
47
        quant_config: Optional[QuantizationConfig] = None,
张大成's avatar
张大成 committed
48
49
50
51
52
    ) -> None:
        super().__init__()
        self.gate_up_proj = MergedColumnParallelLinear(
            hidden_size, [intermediate_size] * 2,
            bias=False,
53
            quant_config=quant_config)
张大成's avatar
张大成 committed
54
55
56
        self.down_proj = RowParallelLinear(intermediate_size,
                                           hidden_size,
                                           bias=False,
57
                                           quant_config=quant_config)
张大成's avatar
张大成 committed
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
        if hidden_act != "silu":
            raise ValueError(f"Unsupported activation: {hidden_act}. "
                             "Only silu is supported for now.")
        self.act_fn = SiluAndMul()

    def forward(self, x):
        gate_up, _ = self.gate_up_proj(x)
        x = self.act_fn(gate_up)
        x, _ = self.down_proj(x)
        return x


class OrionAttention(nn.Module):

    def __init__(
        self,
        hidden_size: int,
        num_heads: int,
        num_kv_heads: int,
        rope_theta: float = 10000,
78
        rope_scaling: Optional[dict[str, Any]] = None,
张大成's avatar
张大成 committed
79
        max_position_embeddings: int = 8192,
80
        cache_config: Optional[CacheConfig] = None,
81
        quant_config: Optional[QuantizationConfig] = None,
82
        prefix: str = "",
张大成's avatar
张大成 committed
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
    ) -> None:
        super().__init__()
        self.hidden_size = hidden_size
        tp_size = get_tensor_model_parallel_world_size()
        self.total_num_heads = num_heads
        assert self.total_num_heads % tp_size == 0
        self.num_heads = self.total_num_heads // tp_size
        self.total_num_kv_heads = num_kv_heads
        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)
        self.head_dim = hidden_size // self.total_num_heads
        self.q_size = self.num_heads * self.head_dim
        self.kv_size = self.num_kv_heads * self.head_dim
        self.scaling = self.head_dim**-0.5
        self.rope_theta = rope_theta
        self.max_position_embeddings = max_position_embeddings

        self.qkv_proj = QKVParallelLinear(
            hidden_size,
            self.head_dim,
            self.total_num_heads,
            self.total_num_kv_heads,
            bias=False,
113
            quant_config=quant_config,
张大成's avatar
张大成 committed
114
115
116
117
118
        )
        self.o_proj = RowParallelLinear(
            self.total_num_heads * self.head_dim,
            hidden_size,
            bias=False,
119
            quant_config=quant_config,
张大成's avatar
张大成 committed
120
121
122
123
124
125
126
127
128
        )

        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,
        )
129
130
131
        self.attn = Attention(self.num_heads,
                              self.head_dim,
                              self.scaling,
132
                              num_kv_heads=self.num_kv_heads,
133
                              cache_config=cache_config,
134
135
                              quant_config=quant_config,
                              prefix=f"{prefix}.attn")
张大成's avatar
张大成 committed
136
137
138
139
140
141
142
143
144

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        qkv, _ = self.qkv_proj(hidden_states)
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
        q, k = self.rotary_emb(positions, q, k)
145
        attn_output = self.attn(q, k, v)
张大成's avatar
张大成 committed
146
147
148
149
150
151
152
153
154
        output, _ = self.o_proj(attn_output)
        return output


class OrionDecoderLayer(nn.Module):

    def __init__(
        self,
        config: PretrainedConfig,
155
        cache_config: Optional[CacheConfig] = None,
156
        quant_config: Optional[QuantizationConfig] = None,
157
        prefix: str = "",
张大成's avatar
张大成 committed
158
159
160
161
162
163
164
165
166
167
168
169
170
171
    ) -> None:
        super().__init__()
        self.hidden_size = config.hidden_size
        rope_theta = getattr(config, "rope_theta", 10000)
        rope_scaling = getattr(config, "rope_scaling", None)
        max_position_embeddings = getattr(config, "max_position_embeddings",
                                          8192)
        self.self_attn = OrionAttention(
            hidden_size=self.hidden_size,
            num_heads=config.num_attention_heads,
            num_kv_heads=config.num_key_value_heads,
            rope_theta=rope_theta,
            rope_scaling=rope_scaling,
            max_position_embeddings=max_position_embeddings,
172
            cache_config=cache_config,
173
            quant_config=quant_config,
174
            prefix=f"{prefix}.self_attn",
张大成's avatar
张大成 committed
175
176
177
178
179
        )
        self.mlp = OrionMLP(
            hidden_size=self.hidden_size,
            intermediate_size=config.intermediate_size,
            hidden_act=config.hidden_act,
180
            quant_config=quant_config,
张大成's avatar
张大成 committed
181
182
183
184
185
186
187
188
189
190
191
        )

        self.input_layernorm = nn.LayerNorm(config.hidden_size,
                                            eps=config.rms_norm_eps)
        self.post_attention_layernorm = nn.LayerNorm(config.hidden_size,
                                                     eps=config.rms_norm_eps)

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
192
    ) -> tuple[torch.Tensor, torch.Tensor]:
张大成's avatar
张大成 committed
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
        # Self Attention
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
        hidden_states = self.self_attn(
            positions=positions,
            hidden_states=hidden_states,
        )

        hidden_states = residual + hidden_states

        # Fully Connected
        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = residual + hidden_states
208
        return hidden_states
张大成's avatar
张大成 committed
209
210


211
@support_torch_compile
张大成's avatar
张大成 committed
212
213
class OrionModel(nn.Module):

214
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
张大成's avatar
张大成 committed
215
        super().__init__()
216
217
218
219
220

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

张大成's avatar
张大成 committed
221
222
223
224
225
226
        self.config = config
        self.vocab_size = config.vocab_size
        self.embed_tokens = VocabParallelEmbedding(
            config.vocab_size,
            config.hidden_size,
        )
227
228
229
        self.start_layer, self.end_layer, self.layers = make_layers(
            config.num_hidden_layers,
            lambda prefix: OrionDecoderLayer(
230
                config, cache_config, quant_config, prefix=prefix),
231
            prefix=f"{prefix}.layers")
张大成's avatar
张大成 committed
232
        self.norm = nn.LayerNorm(config.hidden_size, eps=config.rms_norm_eps)
233
        self.make_empty_intermediate_tensors = (
234
235
236
            make_empty_intermediate_tensors_factory([
                "hidden_states",
            ], config.hidden_size))
张大成's avatar
张大成 committed
237

238
239
240
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embed_tokens(input_ids)

张大成's avatar
张大成 committed
241
242
243
244
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
245
        intermediate_tensors: Optional[IntermediateTensors],
246
        inputs_embeds: Optional[torch.Tensor] = None,
247
248
    ) -> Union[torch.Tensor, IntermediateTensors]:
        if get_pp_group().is_first_rank:
249
250
251
252
            if inputs_embeds is not None:
                hidden_states = inputs_embeds
            else:
                hidden_states = self.get_input_embeddings(input_ids)
253
        else:
254
            assert intermediate_tensors is not None
255
            hidden_states = intermediate_tensors["hidden_states"]
256
        for layer in islice(self.layers, self.start_layer, self.end_layer):
257
            hidden_states = layer(positions, hidden_states)
258
259
260
261
        if not get_pp_group().is_last_rank:
            return IntermediateTensors({
                "hidden_states": hidden_states,
            })
张大成's avatar
张大成 committed
262
263
264
        hidden_states = self.norm(hidden_states)
        return hidden_states

265
266
    def load_weights(self, weights: Iterable[tuple[str,
                                                   torch.Tensor]]) -> set[str]:
267
268
269
270
271
272
273
274
275
        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())
276
        loaded_params: set[str] = set()
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
        for name, loaded_weight in weights:
            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]
                weight_loader = getattr(param, "weight_loader",
                                        default_weight_loader)
                weight_loader(param, loaded_weight)
            loaded_params.add(name)
        return loaded_params

张大成's avatar
张大成 committed
304

305
class OrionForCausalLM(nn.Module, SupportsPP):
张大成's avatar
张大成 committed
306

307
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
张大成's avatar
张大成 committed
308
        super().__init__()
309
310
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
张大成's avatar
张大成 committed
311
        self.config = config
312
        self.quant_config = quant_config
313
314
        self.model = OrionModel(vllm_config=vllm_config,
                                prefix=maybe_prefix(prefix, "model"))
315
316
        self.lm_head = ParallelLMHead(config.vocab_size,
                                      config.hidden_size,
317
318
                                      quant_config=quant_config,
                                      prefix=maybe_prefix(prefix, "lm_head"))
319
320
        if self.config.tie_word_embeddings:
            self.lm_head.weight = self.model.embed_tokens.weight
321
        self.logits_processor = LogitsProcessor(config.vocab_size)
322
323
        self.make_empty_intermediate_tensors = (
            self.model.make_empty_intermediate_tensors)
张大成's avatar
张大成 committed
324

325
326
327
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.model.get_input_embeddings(input_ids)

张大成's avatar
张大成 committed
328
329
330
331
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
332
        intermediate_tensors: Optional[IntermediateTensors] = None,
333
        inputs_embeds: Optional[torch.Tensor] = None,
334
    ) -> Union[torch.Tensor, IntermediateTensors]:
335
        hidden_states = self.model(input_ids, positions, intermediate_tensors,
336
                                   inputs_embeds)
张大成's avatar
张大成 committed
337
338
        return hidden_states

339
340
341
342
343
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
344
        logits = self.logits_processor(self.lm_head, hidden_states,
345
346
347
                                       sampling_metadata)
        return logits

348
349
    def load_weights(self, weights: Iterable[tuple[str,
                                                   torch.Tensor]]) -> set[str]:
350
        loader = AutoWeightsLoader(self)
351
        return loader.load_weights(weights)