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

ywfang's avatar
ywfang committed
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
# Adapted from
# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/llama/modeling_llama.py
# Copyright 2024 The ModelBest team.
# Copyright 2023 The vLLM team.
# 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.
"""Inference-only MiniCPM3 model compatible with HuggingFace weights."""
27

28
from typing import Any
ywfang's avatar
ywfang committed
29
30
31

import torch
from torch import nn
32
from transformers import PretrainedConfig
ywfang's avatar
ywfang committed
33

34
from vllm.attention import Attention
35
from vllm.config import CacheConfig, VllmConfig
ywfang's avatar
ywfang committed
36
37
from vllm.distributed import get_tensor_model_parallel_world_size
from vllm.model_executor.layers.layernorm import RMSNorm
38
39
40
41
42
from vllm.model_executor.layers.linear import (
    ColumnParallelLinear,
    ReplicatedLinear,
    RowParallelLinear,
)
43
from vllm.model_executor.layers.quantization import QuantizationConfig
ywfang's avatar
ywfang committed
44
from vllm.model_executor.layers.rotary_embedding import get_rope
45
46
47
48
49
from vllm.model_executor.models.minicpm import (
    MiniCPMDecoderLayer,
    MiniCPMForCausalLM,
    MiniCPMModel,
)
ywfang's avatar
ywfang committed
50

51
from .utils import make_layers
52

ywfang's avatar
ywfang committed
53
54
55
56

class MiniCPM3Attention(nn.Module):
    def __init__(
        self,
57
        config: PretrainedConfig,
ywfang's avatar
ywfang committed
58
59
60
61
62
63
64
65
        hidden_size: int,
        num_heads: int,
        qk_nope_head_dim: int,
        qk_rope_head_dim: int,
        v_head_dim: int,
        q_lora_rank: int,
        kv_lora_rank: int,
        rope_theta: float = 10000,
66
        rope_scaling: dict[str, Any] | None = None,
ywfang's avatar
ywfang committed
67
        max_position_embeddings: int = 8192,
68
69
        cache_config: CacheConfig | None = None,
        quant_config: QuantizationConfig | None = None,
70
        prefix: str = "",
ywfang's avatar
ywfang committed
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
    ) -> None:
        super().__init__()
        self.hidden_size = hidden_size
        self.qk_nope_head_dim = qk_nope_head_dim
        self.qk_rope_head_dim = qk_rope_head_dim
        self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
        self.v_head_dim = v_head_dim
        self.q_lora_rank = q_lora_rank
        self.kv_lora_rank = kv_lora_rank
        self.num_heads = num_heads

        tp_size = get_tensor_model_parallel_world_size()
        assert self.num_heads % tp_size == 0
        self.num_local_heads = num_heads // tp_size

        self.scaling = self.qk_head_dim**-0.5
        self.rope_theta = rope_theta
        self.max_position_embeddings = max_position_embeddings

90
91
92
        self.q_a_proj = ReplicatedLinear(
            self.hidden_size, self.q_lora_rank, bias=False, quant_config=quant_config
        )
ywfang's avatar
ywfang committed
93
        self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps)
94
95
96
97
98
        self.q_b_proj = ColumnParallelLinear(
            q_lora_rank,
            self.num_heads * self.qk_head_dim,
            bias=False,
            quant_config=quant_config,
99
            prefix=f"{prefix}.q_b_proj",
100
101
102
103
104
105
106
        )

        self.kv_a_proj_with_mqa = ReplicatedLinear(
            self.hidden_size,
            self.kv_lora_rank + self.qk_rope_head_dim,
            bias=False,
            quant_config=quant_config,
107
            prefix=f"{prefix}.kv_a_proj_with_mqa",
108
109
        )
        self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps)
ywfang's avatar
ywfang committed
110
111
112
113
        self.kv_b_proj = ColumnParallelLinear(
            self.kv_lora_rank,
            self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
            bias=False,
114
            quant_config=quant_config,
115
            prefix=f"{prefix}.kv_b_proj",
116
        )
ywfang's avatar
ywfang committed
117
        # O projection.
118
119
120
121
122
        self.o_proj = RowParallelLinear(
            self.num_heads * self.v_head_dim,
            self.hidden_size,
            bias=False,
            quant_config=quant_config,
123
            prefix=f"{prefix}.o_proj",
124
        )
ywfang's avatar
ywfang committed
125
126
127
128
129
130
131
132

        self.rotary_emb = get_rope(
            self.qk_rope_head_dim,
            rotary_dim=self.qk_rope_head_dim,
            max_position=max_position_embeddings,
            base=rope_theta,
            rope_scaling=rope_scaling,
        )
133
134
135
136
137
138
139
140
141
        self.attn = Attention(
            self.num_local_heads,
            self.qk_head_dim,
            self.scaling,
            num_kv_heads=self.num_local_heads,
            cache_config=cache_config,
            quant_config=quant_config,
            prefix=f"{prefix}.attn",
        )
ywfang's avatar
ywfang committed
142
143
144
145
146
147
148
149
150
151

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        q, _ = self.q_a_proj(hidden_states)
        q = self.q_a_layernorm(q)
        q, _ = self.q_b_proj(q)
        q = q.view(-1, self.num_local_heads, self.qk_head_dim)
152
        _, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
ywfang's avatar
ywfang committed
153
        latent_cache, _ = self.kv_a_proj_with_mqa(hidden_states)
154
        kv_a, _ = latent_cache.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
ywfang's avatar
ywfang committed
155
156
157
        latent_cache = latent_cache.unsqueeze(1)
        kv_a = self.kv_a_layernorm(kv_a.contiguous())
        kv, _ = self.kv_b_proj(kv_a)
158
        kv = kv.view(-1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim)
ywfang's avatar
ywfang committed
159
160
        k_nope, v = kv.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)

161
        k_pe = latent_cache[:, :, self.kv_lora_rank :]
ywfang's avatar
ywfang committed
162
163
164
165

        q_pe, k_pe = self.rotary_emb(
            positions,
            q_pe.reshape(-1, self.num_local_heads * self.qk_rope_head_dim),
166
167
            k_pe.reshape(-1, self.qk_rope_head_dim),
        )
ywfang's avatar
ywfang committed
168
169
170
        q_pe = q_pe.view(-1, self.num_local_heads, self.qk_rope_head_dim)
        k_pe = k_pe.view(-1, 1, self.qk_rope_head_dim)

171
        q[..., self.qk_nope_head_dim :] = q_pe
ywfang's avatar
ywfang committed
172
173
174

        k = torch.empty_like(q)

175
176
        k[..., : self.qk_nope_head_dim] = k_nope
        k[..., self.qk_nope_head_dim :] = k_pe
ywfang's avatar
ywfang committed
177
178
179
180

        q = q.reshape(-1, self.num_local_heads * self.qk_head_dim)
        k = k.view(-1, self.num_local_heads * self.qk_head_dim)
        v = torch.nn.functional.pad(
181
182
            v, [0, self.qk_head_dim - self.v_head_dim], value=0
        ).view(-1, self.num_local_heads * self.qk_head_dim)
ywfang's avatar
ywfang committed
183

184
        attn_output = self.attn(q, k, v)
185
186
187
        attn_output = attn_output.view(-1, self.num_local_heads, self.qk_head_dim)[
            ..., : self.v_head_dim
        ].reshape(-1, self.num_local_heads * self.v_head_dim)
ywfang's avatar
ywfang committed
188
189
190
191
192
193
194

        output, _ = self.o_proj(attn_output)
        return output


class MiniCPM3DecoderLayer(MiniCPMDecoderLayer):
    def _init_attn_block(self):
195
196
197
        self.input_layernorm = RMSNorm(
            self.config.hidden_size, eps=self.config.rms_norm_eps
        )
ywfang's avatar
ywfang committed
198
199
200
201
202
203
204
205
206
207
208
209
210
211
        self.self_attn = MiniCPM3Attention(
            config=self.config,
            hidden_size=self.hidden_size,
            num_heads=self.config.num_attention_heads,
            qk_nope_head_dim=self.config.qk_nope_head_dim,
            qk_rope_head_dim=self.config.qk_rope_head_dim,
            v_head_dim=self.config.v_head_dim,
            q_lora_rank=self.config.q_lora_rank,
            kv_lora_rank=self.config.kv_lora_rank,
            rope_theta=self.rope_theta,
            rope_scaling=self.rope_scaling,
            max_position_embeddings=self.max_position_embeddings,
            cache_config=self.cache_config,
            quant_config=self.quant_config,
212
            prefix=f"{self.prefix}.self_attn",
ywfang's avatar
ywfang committed
213
214
215
216
        )


class MiniCPM3Model(MiniCPMModel):
217
218
219
220
    def _init_layers(
        self,
        prefix: str,
        config: PretrainedConfig,
221
222
        cache_config: CacheConfig | None,
        quant_config: QuantizationConfig | None,
223
224
225
    ):
        self.start_layer, self.end_layer, self.layers = make_layers(
            config.num_hidden_layers,
226
            lambda prefix: MiniCPM3DecoderLayer(
227
228
229
230
                config, cache_config, quant_config, prefix=prefix
            ),
            prefix=f"{prefix}.layers",
        )
ywfang's avatar
ywfang committed
231
232
233


class MiniCPM3ForCausalLM(MiniCPMForCausalLM):
234
235
236
237
238
239
240
    packed_modules_mapping = {
        "gate_up_proj": [
            "gate_proj",
            "up_proj",
        ],
    }

241
    def _init_model(self, *, vllm_config: VllmConfig, prefix: str = ""):
242
        return MiniCPM3Model(vllm_config=vllm_config, prefix=prefix)