minicpm3.py 8.37 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

ywfang's avatar
ywfang committed
28
29
import torch
from torch import nn
30
from transformers import PretrainedConfig
ywfang's avatar
ywfang committed
31

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

49
from .utils import make_layers
50

ywfang's avatar
ywfang committed
51
52
53
54

class MiniCPM3Attention(nn.Module):
    def __init__(
        self,
55
        config: PretrainedConfig,
ywfang's avatar
ywfang committed
56
57
58
59
60
61
62
63
        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,
        max_position_embeddings: int = 8192,
64
65
        cache_config: CacheConfig | None = None,
        quant_config: QuantizationConfig | None = None,
66
        prefix: str = "",
ywfang's avatar
ywfang committed
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
    ) -> 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.max_position_embeddings = max_position_embeddings

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

        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,
102
            prefix=f"{prefix}.kv_a_proj_with_mqa",
103
104
        )
        self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps)
ywfang's avatar
ywfang committed
105
106
107
108
        self.kv_b_proj = ColumnParallelLinear(
            self.kv_lora_rank,
            self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
            bias=False,
109
            quant_config=quant_config,
110
            prefix=f"{prefix}.kv_b_proj",
111
        )
ywfang's avatar
ywfang committed
112
        # O projection.
113
114
115
116
117
        self.o_proj = RowParallelLinear(
            self.num_heads * self.v_head_dim,
            self.hidden_size,
            bias=False,
            quant_config=quant_config,
118
            prefix=f"{prefix}.o_proj",
119
        )
ywfang's avatar
ywfang committed
120
121
122
123

        self.rotary_emb = get_rope(
            self.qk_rope_head_dim,
            max_position=max_position_embeddings,
124
            rope_parameters=config.rope_parameters,
ywfang's avatar
ywfang committed
125
        )
126
127
128
129
130
131
132
133
134
        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
135
136
137
138
139
140
141
142
143
144

    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)
145
        _, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
ywfang's avatar
ywfang committed
146
        latent_cache, _ = self.kv_a_proj_with_mqa(hidden_states)
147
        kv_a, _ = latent_cache.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
ywfang's avatar
ywfang committed
148
149
150
        latent_cache = latent_cache.unsqueeze(1)
        kv_a = self.kv_a_layernorm(kv_a.contiguous())
        kv, _ = self.kv_b_proj(kv_a)
151
        kv = kv.view(-1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim)
ywfang's avatar
ywfang committed
152
153
        k_nope, v = kv.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)

154
        k_pe = latent_cache[:, :, self.kv_lora_rank :]
ywfang's avatar
ywfang committed
155
156
157
158

        q_pe, k_pe = self.rotary_emb(
            positions,
            q_pe.reshape(-1, self.num_local_heads * self.qk_rope_head_dim),
159
160
            k_pe.reshape(-1, self.qk_rope_head_dim),
        )
ywfang's avatar
ywfang committed
161
162
163
        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)

164
        q[..., self.qk_nope_head_dim :] = q_pe
ywfang's avatar
ywfang committed
165
166
167

        k = torch.empty_like(q)

168
169
        k[..., : self.qk_nope_head_dim] = k_nope
        k[..., self.qk_nope_head_dim :] = k_pe
ywfang's avatar
ywfang committed
170
171
172
173

        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(
174
175
            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
176

177
        attn_output = self.attn(q, k, v)
178
179
180
        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
181
182
183
184
185
186
187

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


class MiniCPM3DecoderLayer(MiniCPMDecoderLayer):
    def _init_attn_block(self):
188
189
190
        self.input_layernorm = RMSNorm(
            self.config.hidden_size, eps=self.config.rms_norm_eps
        )
ywfang's avatar
ywfang committed
191
192
193
194
195
196
197
198
199
200
201
202
        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,
            max_position_embeddings=self.max_position_embeddings,
            cache_config=self.cache_config,
            quant_config=self.quant_config,
203
            prefix=f"{self.prefix}.self_attn",
ywfang's avatar
ywfang committed
204
205
206
207
        )


class MiniCPM3Model(MiniCPMModel):
208
209
210
211
    def _init_layers(
        self,
        prefix: str,
        config: PretrainedConfig,
212
213
        cache_config: CacheConfig | None,
        quant_config: QuantizationConfig | None,
214
215
216
    ):
        self.start_layer, self.end_layer, self.layers = make_layers(
            config.num_hidden_layers,
217
            lambda prefix: MiniCPM3DecoderLayer(
218
219
220
221
                config, cache_config, quant_config, prefix=prefix
            ),
            prefix=f"{prefix}.layers",
        )
ywfang's avatar
ywfang committed
222
223
224


class MiniCPM3ForCausalLM(MiniCPMForCausalLM):
225
226
227
228
229
230
231
    packed_modules_mapping = {
        "gate_up_proj": [
            "gate_proj",
            "up_proj",
        ],
    }

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