orion.py 13.8 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0

张大成's avatar
张大成 committed
3
4
5
6
7
# 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."""
8
from typing import Any, Dict, Iterable, Optional, Set, Tuple, Union
张大成's avatar
张大成 committed
9
10
11
12
13

import torch
from torch import nn
from transformers import PretrainedConfig

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

31
from .interfaces import SupportsPP
32
from .utils import (AutoWeightsLoader, is_pp_missing_parameter,
33
34
                    make_empty_intermediate_tensors_factory, make_layers,
                    maybe_prefix)
35

张大成's avatar
张大成 committed
36
37
38
39
40
41
42
43

class OrionMLP(nn.Module):

    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
        hidden_act: str,
44
        quant_config: Optional[QuantizationConfig] = None,
张大成's avatar
张大成 committed
45
46
47
48
49
    ) -> None:
        super().__init__()
        self.gate_up_proj = MergedColumnParallelLinear(
            hidden_size, [intermediate_size] * 2,
            bias=False,
50
            quant_config=quant_config)
张大成's avatar
张大成 committed
51
52
53
        self.down_proj = RowParallelLinear(intermediate_size,
                                           hidden_size,
                                           bias=False,
54
                                           quant_config=quant_config)
张大成's avatar
张大成 committed
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
        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,
        rope_scaling: Optional[Dict[str, Any]] = None,
        max_position_embeddings: int = 8192,
77
        cache_config: Optional[CacheConfig] = None,
78
        quant_config: Optional[QuantizationConfig] = None,
79
        prefix: str = "",
张大成's avatar
张大成 committed
80
81
82
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
    ) -> 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,
110
            quant_config=quant_config,
张大成's avatar
张大成 committed
111
112
113
114
115
        )
        self.o_proj = RowParallelLinear(
            self.total_num_heads * self.head_dim,
            hidden_size,
            bias=False,
116
            quant_config=quant_config,
张大成's avatar
张大成 committed
117
118
119
120
121
122
123
124
125
        )

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

    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)
142
        attn_output = self.attn(q, k, v)
张大成's avatar
张大成 committed
143
144
145
146
147
148
149
150
151
        output, _ = self.o_proj(attn_output)
        return output


class OrionDecoderLayer(nn.Module):

    def __init__(
        self,
        config: PretrainedConfig,
152
        cache_config: Optional[CacheConfig] = None,
153
        quant_config: Optional[QuantizationConfig] = None,
154
        prefix: str = "",
张大成's avatar
张大成 committed
155
156
157
158
159
160
161
162
163
164
165
166
167
168
    ) -> 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,
169
            cache_config=cache_config,
170
            quant_config=quant_config,
171
            prefix=f"{prefix}.self_attn",
张大成's avatar
张大成 committed
172
173
174
175
176
        )
        self.mlp = OrionMLP(
            hidden_size=self.hidden_size,
            intermediate_size=config.intermediate_size,
            hidden_act=config.hidden_act,
177
            quant_config=quant_config,
张大成's avatar
张大成 committed
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
        )

        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,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        # 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
205
        return hidden_states
张大成's avatar
张大成 committed
206
207


208
@support_torch_compile
张大成's avatar
张大成 committed
209
210
class OrionModel(nn.Module):

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

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

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

235
236
237
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embed_tokens(input_ids)

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

262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
        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())
        loaded_params: Set[str] = set()
        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
301

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

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

321
322
323
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.model.get_input_embeddings(input_ids)

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

335
336
337
338
339
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
340
        logits = self.logits_processor(self.lm_head, hidden_states,
341
342
343
                                       sampling_metadata)
        return logits

344
345
    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
346
347
348
349
        loader = AutoWeightsLoader(
            self,
            skip_prefixes=([
                "rotary_emb.inv_freq",
张大成's avatar
张大成 committed
350
351
                # Models trained using ColossalAI may include these tensors in
                # the checkpoint. Skip them.
352
353
354
355
356
                "rotary_emb.cos_cached",
                "rotary_emb.sin_cached"
            ]),
        )
        return loader.load_weights(weights)