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

wangding zeng's avatar
wangding zeng committed
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
# Adapted from
# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/llama/modeling_llama.py
# Copyright 2023 The vLLM team.
# Copyright 2023 DeepSeek-AI 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.
25
"""Inference-only DeepseekV2/DeepseekV3 model."""
26
27
import typing
from collections.abc import Callable, Iterable
28
from itertools import islice
29
from typing import Any, Optional, Union
wangding zeng's avatar
wangding zeng committed
30

31
32
import os
import re
wangding zeng's avatar
wangding zeng committed
33
34
import torch
from torch import nn
35
from transformers import DeepseekV2Config, DeepseekV3Config
wangding zeng's avatar
wangding zeng committed
36

37
from vllm.attention import Attention
38
from vllm.compilation.decorators import support_torch_compile
39
40
41
42
from vllm.config import (CacheConfig, ModelConfig, VllmConfig,
                         get_current_vllm_config)
from vllm.distributed import (get_ep_group, get_pp_group,
                              get_tensor_model_parallel_world_size)
wangding zeng's avatar
wangding zeng committed
43
from vllm.model_executor.layers.activation import SiluAndMul
44
from vllm.model_executor.layers.fused_moe import FusedMoE
wangding zeng's avatar
wangding zeng committed
45
46
47
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
                                               MergedColumnParallelLinear,
48
                                               MergedReplicatedLinear,
wangding zeng's avatar
wangding zeng committed
49
50
51
                                               ReplicatedLinear,
                                               RowParallelLinear)
from vllm.model_executor.layers.logits_processor import LogitsProcessor
52
from vllm.model_executor.layers.quantization import QuantizationConfig
wangding zeng's avatar
wangding zeng committed
53
54
55
from vllm.model_executor.layers.rotary_embedding import get_rope
from vllm.model_executor.layers.vocab_parallel_embedding import (
    ParallelLMHead, VocabParallelEmbedding)
56
57
from vllm.model_executor.model_loader.weight_utils import (
    default_weight_loader, maybe_remap_kv_scale_name)
wangding zeng's avatar
wangding zeng committed
58
from vllm.model_executor.sampling_metadata import SamplingMetadata
59
from vllm.sequence import IntermediateTensors
wangding zeng's avatar
wangding zeng committed
60

61
from .interfaces import MixtureOfExperts, SupportsLoRA, SupportsPP
62
from .utils import (PPMissingLayer, is_pp_missing_parameter,
63
64
                    make_empty_intermediate_tensors_factory, make_layers,
                    maybe_prefix)
65
from vllm import _custom_ops as ops
66

wangding zeng's avatar
wangding zeng committed
67
68
69
70
71
72
73
74
75
76

class DeepseekV2MLP(nn.Module):

    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
        hidden_act: str,
        quant_config: Optional[QuantizationConfig] = None,
        reduce_results: bool = True,
77
        prefix: str = "",
wangding zeng's avatar
wangding zeng committed
78
79
80
81
82
    ) -> None:
        super().__init__()
        self.gate_up_proj = MergedColumnParallelLinear(
            hidden_size, [intermediate_size] * 2,
            bias=False,
83
84
            quant_config=quant_config,
            prefix=f"{prefix}.gate_up_proj")
wangding zeng's avatar
wangding zeng committed
85
86
87
88
        self.down_proj = RowParallelLinear(intermediate_size,
                                           hidden_size,
                                           bias=False,
                                           quant_config=quant_config,
89
90
                                           reduce_results=reduce_results,
                                           prefix=f"{prefix}.down_proj")
wangding zeng's avatar
wangding zeng committed
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
        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 DeepseekV2MoE(nn.Module):

    def __init__(
        self,
107
        config: Union[DeepseekV2Config, DeepseekV3Config],
wangding zeng's avatar
wangding zeng committed
108
        quant_config: Optional[QuantizationConfig] = None,
109
        prefix: str = "",
110
        enable_eplb: bool = False,
wangding zeng's avatar
wangding zeng committed
111
112
113
114
    ):
        super().__init__()
        self.tp_size = get_tensor_model_parallel_world_size()
        self.routed_scaling_factor = config.routed_scaling_factor
115
116
117
118
119
120

        self.ep_group = get_ep_group().device_group
        self.ep_rank = self.ep_group.rank()
        self.ep_size = self.ep_group.size()
        self.n_routed_experts: int = config.n_routed_experts
        self.n_shared_experts: int = config.n_shared_experts
121
122
123
124
125

        if config.hidden_act != "silu":
            raise ValueError(f"Unsupported activation: {config.hidden_act}. "
                             "Only silu is supported for now.")

wangding zeng's avatar
wangding zeng committed
126
        self.gate = ReplicatedLinear(config.hidden_size,
127
                                     config.n_routed_experts,
wangding zeng's avatar
wangding zeng committed
128
                                     bias=False,
129
130
                                     quant_config=None,
                                     prefix=f"{prefix}.gate")
131
132
        if config.topk_method == "noaux_tc":
            self.gate.e_score_correction_bias = nn.Parameter(
133
                torch.empty(config.n_routed_experts, dtype=torch.float32))
134
135
136
        else:
            self.gate.e_score_correction_bias = None

137
138
        # Load balancing settings.
        vllm_config = get_current_vllm_config()
139
        eplb_config = vllm_config.parallel_config.eplb_config
140
141
        self.enable_eplb = enable_eplb

142
        self.n_redundant_experts = eplb_config.num_redundant_experts
143
144
145
146
147
148
149
150
151
152
        self.n_logical_experts = self.n_routed_experts
        self.n_physical_experts = (self.n_logical_experts +
                                   self.n_redundant_experts)
        self.n_local_physical_experts = self.n_physical_experts // self.ep_size

        self.physical_expert_start = (self.ep_rank *
                                      self.n_local_physical_experts)
        self.physical_expert_end = (self.physical_expert_start +
                                    self.n_local_physical_experts)

153
154
155
156
157
158
159
160
161
162
163
164
165
        self.experts = FusedMoE(
            num_experts=config.n_routed_experts,
            top_k=config.num_experts_per_tok,
            hidden_size=config.hidden_size,
            intermediate_size=config.moe_intermediate_size,
            reduce_results=False,
            renormalize=config.norm_topk_prob,
            quant_config=quant_config,
            use_grouped_topk=True,
            num_expert_group=config.n_group,
            topk_group=config.topk_group,
            prefix=f"{prefix}.experts",
            scoring_func=config.scoring_func,
166
            routed_scaling_factor=self.routed_scaling_factor,
167
168
169
            e_score_correction_bias=self.gate.e_score_correction_bias,
            enable_eplb=self.enable_eplb,
            num_redundant_experts=self.n_redundant_experts)
170

wangding zeng's avatar
wangding zeng committed
171
172
173
174
175
176
177
178
        if config.n_shared_experts is not None:
            intermediate_size = (config.moe_intermediate_size *
                                 config.n_shared_experts)
            self.shared_experts = DeepseekV2MLP(
                hidden_size=config.hidden_size,
                intermediate_size=intermediate_size,
                hidden_act=config.hidden_act,
                quant_config=quant_config,
179
180
                reduce_results=self.experts.must_reduce_shared_expert_outputs(
                ),
181
                prefix=f"{prefix}.shared_experts",
wangding zeng's avatar
wangding zeng committed
182
183
184
185
186
            )

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        num_tokens, hidden_dim = hidden_states.shape
        hidden_states = hidden_states.view(-1, hidden_dim)
187
        if self.n_shared_experts is not None:
wangding zeng's avatar
wangding zeng committed
188
189
190
            shared_output = self.shared_experts(hidden_states)
        # router_logits: (num_tokens, n_experts)
        router_logits, _ = self.gate(hidden_states)
191

192
193
194
195
196
        if hidden_states.dtype != torch.float16:
            final_hidden_states = self.experts(
                hidden_states=hidden_states,
                router_logits=router_logits) * self.routed_scaling_factor
        else:
197
198
            # Fix FP16 overflow
            # See DeepseekV2DecoderLayer for more details.
199
200
            final_hidden_states = self.experts(hidden_states=hidden_states,
                                               router_logits=router_logits)
201
        if shared_output is not None:
202
203
204
            if hidden_states.dtype != torch.float16:
                final_hidden_states = final_hidden_states + shared_output
            else:
205
206
                # Fix FP16 overflow
                # See DeepseekV2DecoderLayer for more details.
207
208
                final_hidden_states = final_hidden_states + shared_output \
                    * (1. / self.routed_scaling_factor)
209

210
        if self.tp_size > 1:
211
212
213
            final_hidden_states = (
                self.experts.maybe_all_reduce_tensor_model_parallel(
                    final_hidden_states))
wangding zeng's avatar
wangding zeng committed
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228

        return final_hidden_states.view(num_tokens, hidden_dim)


def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float:
    import math
    if scale <= 1:
        return 1.0
    return 0.1 * mscale * math.log(scale) + 1.0


class DeepseekV2Attention(nn.Module):

    def __init__(
        self,
229
        config: Union[DeepseekV2Config, DeepseekV3Config],
wangding zeng's avatar
wangding zeng committed
230
231
232
233
234
235
236
237
        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,
238
        rope_scaling: Optional[dict[str, Any]] = None,
wangding zeng's avatar
wangding zeng committed
239
240
241
        max_position_embeddings: int = 8192,
        cache_config: Optional[CacheConfig] = None,
        quant_config: Optional[QuantizationConfig] = None,
242
        prefix: str = "",
wangding zeng's avatar
wangding zeng committed
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
    ) -> 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 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

        if self.q_lora_rank is not None:
            self.q_a_proj = ReplicatedLinear(self.hidden_size,
                                             self.q_lora_rank,
                                             bias=False,
264
265
                                             quant_config=quant_config,
                                             prefix=f"{prefix}.q_a_proj")
wangding zeng's avatar
wangding zeng committed
266
267
268
269
270
271
            self.q_a_layernorm = RMSNorm(self.q_lora_rank,
                                         eps=config.rms_norm_eps)
            self.q_b_proj = ColumnParallelLinear(q_lora_rank,
                                                 self.num_heads *
                                                 self.qk_head_dim,
                                                 bias=False,
272
273
                                                 quant_config=quant_config,
                                                 prefix=f"{prefix}.q_b_proj")
wangding zeng's avatar
wangding zeng committed
274
275
276
277
278
        else:
            self.q_proj = ColumnParallelLinear(self.hidden_size,
                                               self.num_heads *
                                               self.qk_head_dim,
                                               bias=False,
279
280
                                               quant_config=quant_config,
                                               prefix=f"{prefix}.q_proj")
wangding zeng's avatar
wangding zeng committed
281

282
283
284
285
286
287
        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,
            prefix=f"{prefix}.kv_a_proj_with_mqa")
wangding zeng's avatar
wangding zeng committed
288
289
290
291
292
293
        self.kv_a_layernorm = RMSNorm(self.kv_lora_rank,
                                      eps=config.rms_norm_eps)
        self.kv_b_proj = ColumnParallelLinear(
            self.kv_lora_rank,
            self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
            bias=False,
294
295
            quant_config=quant_config,
            prefix=f"{prefix}.kv_b_proj")
wangding zeng's avatar
wangding zeng committed
296
297
298
299
        # O projection.
        self.o_proj = RowParallelLinear(self.num_heads * self.v_head_dim,
                                        self.hidden_size,
                                        bias=False,
300
301
                                        quant_config=quant_config,
                                        prefix=f"{prefix}.o_proj")
302
303
        if rope_scaling:
            rope_scaling["rope_type"] = 'deepseek_yarn'
304

wangding zeng's avatar
wangding zeng committed
305
306
307
308
309
310
311
312
313
314
315
316
317
318
        self.rotary_emb = get_rope(qk_rope_head_dim,
                                   rotary_dim=qk_rope_head_dim,
                                   max_position=max_position_embeddings,
                                   base=rope_theta,
                                   rope_scaling=rope_scaling,
                                   is_neox_style=False)

        if rope_scaling:
            mscale_all_dim = rope_scaling.get("mscale_all_dim", False)
            scaling_factor = rope_scaling["factor"]
            mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim))
            self.scaling = self.scaling * mscale * mscale

        self.attn = Attention(self.num_local_heads,
319
                              self.qk_head_dim,
wangding zeng's avatar
wangding zeng committed
320
321
322
                              self.scaling,
                              num_kv_heads=self.num_local_heads,
                              cache_config=cache_config,
323
324
                              quant_config=quant_config,
                              prefix=f"{prefix}.attn")
wangding zeng's avatar
wangding zeng committed
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        if self.q_lora_rank is not None:
            q = self.q_a_proj(hidden_states)[0]
            q = self.q_a_layernorm(q)
            q = self.q_b_proj(q)[0].view(-1, self.num_local_heads,
                                         self.qk_head_dim)
        else:
            q = self.q_proj(hidden_states)[0].view(-1, self.num_local_heads,
                                                   self.qk_head_dim)
        q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim],
                               dim=-1)
        latent_cache = self.kv_a_proj_with_mqa(hidden_states)[0]
        kv_a, _ = latent_cache.split(
            [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
        latent_cache = latent_cache.unsqueeze(1)
345
        kv_a = self.kv_a_layernorm(kv_a)
wangding zeng's avatar
wangding zeng committed
346
347
348
349
350
        kv = self.kv_b_proj(kv_a)[0]
        kv = kv.view(-1, self.num_local_heads,
                     self.qk_nope_head_dim + self.v_head_dim)
        k_nope, v = kv.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
        k_pe = latent_cache[:, :, self.kv_lora_rank:]
351

wangding zeng's avatar
wangding zeng committed
352
        q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe)
353

wangding zeng's avatar
wangding zeng committed
354
355
356
357
        q[..., self.qk_nope_head_dim:] = q_pe
        k = torch.empty_like(q)
        k[..., :self.qk_nope_head_dim] = k_nope
        k[..., self.qk_nope_head_dim:] = k_pe
358
359
360
361
        # padding value to qk_head_dim for alignment
        v = torch.nn.functional.pad(
            v, [0, self.qk_head_dim - self.v_head_dim],
            value=0).view(-1, self.num_local_heads * self.qk_head_dim)
362
        attn_output = self.attn(q, k, v)
wangding zeng's avatar
wangding zeng committed
363
        attn_output = attn_output.view(
364
365
            -1, self.num_local_heads,
            self.qk_head_dim)[..., :self.v_head_dim].reshape(
wangding zeng's avatar
wangding zeng committed
366
367
368
369
370
                -1, self.num_local_heads * self.v_head_dim)
        output, _ = self.o_proj(attn_output)
        return output


371
372
373
374
375
376
377
378
379
380
class DeepseekV2MLAAttention(nn.Module):
    """
    Main reference: DeepseekV2 paper, and FlashInfer Implementation
    (https://arxiv.org/abs/2405.04434 and https://github.com/flashinfer-ai/flashinfer/pull/551).
    
    For more info see MLACommonImpl in: vllm/attention/backends/mla/utils.py
    """

    def __init__(
        self,
381
        config: Union[DeepseekV2Config, DeepseekV3Config],
382
383
384
385
386
387
388
389
        hidden_size: int,
        num_heads: int,
        qk_nope_head_dim: int,
        qk_rope_head_dim: int,
        v_head_dim: int,
        q_lora_rank: Optional[int],
        kv_lora_rank: int,
        rope_theta: float = 10000,
390
        rope_scaling: Optional[dict[str, Any]] = None,
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
        max_position_embeddings: int = 8192,
        cache_config: Optional[CacheConfig] = None,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ) -> 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 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

        if self.q_lora_rank is not None:
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
            self.fused_qkv_a_proj = MergedReplicatedLinear(
                self.hidden_size,
                [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
                bias=False,
                quant_config=quant_config,
                prefix=f"{prefix}.fused_qkv_a_proj")
        else:
            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,
                prefix=f"{prefix}.kv_a_proj_with_mqa")

        if self.q_lora_rank is not None:
431
432
            self.q_a_layernorm = RMSNorm(self.q_lora_rank,
                                         eps=config.rms_norm_eps)
433
            self.q_b_proj = ColumnParallelLinear(self.q_lora_rank,
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
                                                 self.num_heads *
                                                 self.qk_head_dim,
                                                 bias=False,
                                                 quant_config=quant_config,
                                                 prefix=f"{prefix}.q_b_proj")
        else:
            self.q_proj = ColumnParallelLinear(self.hidden_size,
                                               self.num_heads *
                                               self.qk_head_dim,
                                               bias=False,
                                               quant_config=quant_config,
                                               prefix=f"{prefix}.q_proj")
        self.kv_a_layernorm = RMSNorm(self.kv_lora_rank,
                                      eps=config.rms_norm_eps)
        self.kv_b_proj = ColumnParallelLinear(
            self.kv_lora_rank,
            self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.kv_b_proj")
        self.o_proj = RowParallelLinear(self.num_heads * self.v_head_dim,
                                        self.hidden_size,
                                        bias=False,
                                        quant_config=quant_config,
                                        prefix=f"{prefix}.o_proj")

460
461
        if rope_scaling:
            rope_scaling["rope_type"] = 'deepseek_yarn'
462
463
464
465
466
467
468
469
470
471
472
473
        self.rotary_emb = get_rope(qk_rope_head_dim,
                                   rotary_dim=qk_rope_head_dim,
                                   max_position=max_position_embeddings,
                                   base=rope_theta,
                                   rope_scaling=rope_scaling,
                                   is_neox_style=False)
        if rope_scaling:
            mscale_all_dim = rope_scaling.get("mscale_all_dim", False)
            scaling_factor = rope_scaling["factor"]
            mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim))
            self.scaling = self.scaling * mscale * mscale

474
475
476
477
478
479
        # In the MLA backend, kv_cache includes both k_c and
        # pe (i.e. decoupled position embeddings). In particular,
        # the concat_and_cache_mla op requires
        #     k_c.size(1) + k_pe.size(1) == kv_cache.size(2)
        # i.e.
        #     kv_lora_rank + qk_rope_head_dim == head_size
480
481
        self.mla_attn = Attention(
            num_heads=self.num_local_heads,
482
            head_size=self.kv_lora_rank + self.qk_rope_head_dim,
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
            scale=self.scaling,
            num_kv_heads=1,
            cache_config=cache_config,
            quant_config=quant_config,
            prefix=f"{prefix}.attn",
            use_mla=True,
            # MLA Args
            q_lora_rank=self.q_lora_rank,
            kv_lora_rank=self.kv_lora_rank,
            qk_nope_head_dim=self.qk_nope_head_dim,
            qk_rope_head_dim=self.qk_rope_head_dim,
            qk_head_dim=self.qk_head_dim,
            v_head_dim=self.v_head_dim,
            kv_b_proj=self.kv_b_proj,
        )

        self.prefix = prefix
        self.debug_layer_idx = int(self.prefix.split(".")[-2])

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
507
508
509
        q_c = None
        kv_lora = None

510
        if self.q_lora_rank is not None:
511
512
513
514
515
            qkv_lora = self.fused_qkv_a_proj(hidden_states)[0]
            q_c, kv_lora = qkv_lora.split(
                [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
                dim=-1,
            )
516
517
            q_c = self.q_a_layernorm(q_c)
            q = self.q_b_proj(q_c)[0]
518
        else:
519
            kv_lora = self.kv_a_proj_with_mqa(hidden_states)[0]
520
            q = self.q_proj(hidden_states)[0]
521
522
523
524

        kv_c, k_pe = kv_lora.split([self.kv_lora_rank, self.qk_rope_head_dim],
                                   dim=-1)
        kv_c_normed = self.kv_a_layernorm(kv_c)
525

526
527
528
529
530
531
532
        q = q.view(-1, self.num_local_heads, self.qk_head_dim)
        # Add head dim of 1 to k_pe
        k_pe = k_pe.unsqueeze(1)

        q[..., self.qk_nope_head_dim:], k_pe = self.rotary_emb(
            positions, q[..., self.qk_nope_head_dim:], k_pe)

533
534
535
536
537
538
539
        attn_out = self.mla_attn(
            q,
            kv_c_normed,
            k_pe,
            output_shape=(hidden_states.shape[0],
                          self.num_local_heads * self.v_head_dim))
        return self.o_proj(attn_out)[0]
540
541


wangding zeng's avatar
wangding zeng committed
542
543
544
545
class DeepseekV2DecoderLayer(nn.Module):

    def __init__(
        self,
546
        config: Union[DeepseekV2Config, DeepseekV3Config],
547
        prefix: str,
548
        model_config: ModelConfig,
wangding zeng's avatar
wangding zeng committed
549
550
        cache_config: Optional[CacheConfig] = None,
        quant_config: Optional[QuantizationConfig] = None,
551
        enable_eplb: bool = False,
wangding zeng's avatar
wangding zeng committed
552
553
554
555
556
557
558
    ) -> 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)
559
560
561
        # DecoderLayers are created with `make_layers` which passes the prefix
        # with the layer's index.
        layer_idx = int(prefix.split(sep='.')[-1])
562
        self.layer_idx = layer_idx
563
564
565
566
567
        if model_config.use_mla:
            attn_cls = DeepseekV2MLAAttention
        else:
            attn_cls = DeepseekV2Attention
        self.self_attn = attn_cls(
wangding zeng's avatar
wangding zeng committed
568
569
570
571
572
573
574
575
576
577
578
579
580
581
            config=config,
            hidden_size=self.hidden_size,
            num_heads=config.num_attention_heads,
            qk_nope_head_dim=config.qk_nope_head_dim,
            qk_rope_head_dim=config.qk_rope_head_dim,
            v_head_dim=config.v_head_dim,
            q_lora_rank=config.q_lora_rank
            if hasattr(config, "q_lora_rank") else None,
            kv_lora_rank=config.kv_lora_rank,
            rope_theta=rope_theta,
            rope_scaling=rope_scaling,
            max_position_embeddings=max_position_embeddings,
            cache_config=cache_config,
            quant_config=quant_config,
582
            prefix=f"{prefix}.self_attn",
wangding zeng's avatar
wangding zeng committed
583
        )
584

wangding zeng's avatar
wangding zeng committed
585
586
587
        if (config.n_routed_experts is not None
                and layer_idx >= config.first_k_dense_replace
                and layer_idx % config.moe_layer_freq == 0):
588
589
590
591
            self.mlp = DeepseekV2MoE(
                config=config,
                quant_config=quant_config,
                prefix=f"{prefix}.mlp",
592
                enable_eplb=enable_eplb,
593
            )
wangding zeng's avatar
wangding zeng committed
594
595
596
597
598
599
        else:
            self.mlp = DeepseekV2MLP(
                hidden_size=config.hidden_size,
                intermediate_size=config.intermediate_size,
                hidden_act=config.hidden_act,
                quant_config=quant_config,
600
                prefix=f"{prefix}.mlp",
wangding zeng's avatar
wangding zeng committed
601
602
603
604
605
            )
        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)
606
        self.routed_scaling_factor = config.routed_scaling_factor
wangding zeng's avatar
wangding zeng committed
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        residual: Optional[torch.Tensor],
    ) -> torch.Tensor:
        # Self Attention
        if residual is None:
            residual = hidden_states
            hidden_states = self.input_layernorm(hidden_states)
        else:
            hidden_states, residual = self.input_layernorm(
                hidden_states, residual)
        hidden_states = self.self_attn(
            positions=positions,
            hidden_states=hidden_states,
        )

626
627
628
629
        if hidden_states.dtype == torch.float16:
            # Fix FP16 overflow
            # We scale both hidden_states and residual before
            # rmsnorm, and rmsnorm result would not affect by scale.
630
            hidden_states *= 1. / self.routed_scaling_factor
631
632
633
634
635
636
            if self.layer_idx == 0:
                # The residual is shared by all layers, we only scale it on
                # first layer.
                residual *= 1. / self.routed_scaling_factor

        # Fully Connected
wangding zeng's avatar
wangding zeng committed
637
638
639
        hidden_states, residual = self.post_attention_layernorm(
            hidden_states, residual)
        hidden_states = self.mlp(hidden_states)
640
641
642
643
644
645
646
647

        if isinstance(self.mlp,
                      DeepseekV2MLP) and hidden_states.dtype == torch.float16:
            # Fix FP16 overflow
            # Scaling the DeepseekV2MLP output, it is the input of
            # input_layernorm of next decoder layer.
            # The scaling of DeepseekV2MOE output would be done in the forward
            # of DeepseekV2MOE
648
            hidden_states *= 1. / self.routed_scaling_factor
649

wangding zeng's avatar
wangding zeng committed
650
651
652
        return hidden_states, residual


653
@support_torch_compile
wangding zeng's avatar
wangding zeng committed
654
655
656
657
class DeepseekV2Model(nn.Module):

    fall_back_to_pt_during_load = False

658
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
wangding zeng's avatar
wangding zeng committed
659
        super().__init__()
660
661

        config = vllm_config.model_config.hf_config
662
        model_config = vllm_config.model_config
663
664
        cache_config = vllm_config.cache_config
        quant_config = vllm_config.quant_config
665
        enable_eplb = vllm_config.parallel_config.enable_eplb
666
        self.config = config
667

wangding zeng's avatar
wangding zeng committed
668
669
        self.vocab_size = config.vocab_size

670
671
672
673
        if get_pp_group().is_first_rank:
            self.embed_tokens = VocabParallelEmbedding(
                config.vocab_size,
                config.hidden_size,
674
675
                quant_config=quant_config,
                prefix=f"{prefix}.embed_tokens")
676
677
678
679
680
681
682
683
        else:
            self.embed_tokens = PPMissingLayer()

        self.start_layer, self.end_layer, self.layers = make_layers(
            config.num_hidden_layers,
            lambda prefix: DeepseekV2DecoderLayer(
                config,
                prefix,
684
                model_config=model_config,
685
686
                cache_config=cache_config,
                quant_config=quant_config,
687
                enable_eplb=enable_eplb,
688
689
690
691
692
693
694
            ),
            prefix=f"{prefix}.layers")

        if get_pp_group().is_last_rank:
            self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        else:
            self.norm = PPMissingLayer()
695
696
697
        self.make_empty_intermediate_tensors = (
            make_empty_intermediate_tensors_factory(
                ["hidden_states", "residual"], config.hidden_size))
wangding zeng's avatar
wangding zeng committed
698

699
700
701
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embed_tokens(input_ids)

wangding zeng's avatar
wangding zeng committed
702
703
704
705
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
706
        intermediate_tensors: Optional[IntermediateTensors],
707
        inputs_embeds: Optional[torch.Tensor] = None,
708
    ) -> Union[torch.Tensor, IntermediateTensors]:
709
        if get_pp_group().is_first_rank:
710
711
712
713
            if inputs_embeds is not None:
                hidden_states = inputs_embeds
            else:
                hidden_states = self.get_input_embeddings(input_ids)
714
715
716
717
718
719
            residual = None
        else:
            assert intermediate_tensors is not None
            hidden_states = intermediate_tensors["hidden_states"]
            residual = intermediate_tensors["residual"]

720
        for layer in islice(self.layers, self.start_layer, self.end_layer):
721
            hidden_states, residual = layer(positions, hidden_states, residual)
722
723
724
725
726
727
728

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

wangding zeng's avatar
wangding zeng committed
729
730
731
732
        hidden_states, _ = self.norm(hidden_states, residual)
        return hidden_states


733
734
class DeepseekV2ForCausalLM(nn.Module, SupportsPP, MixtureOfExperts,
                            SupportsLoRA):
735
736
737
    packed_modules_mapping = {
        "gate_up_proj": ["gate_proj", "up_proj"],
    }
wangding zeng's avatar
wangding zeng committed
738

739
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
wangding zeng's avatar
wangding zeng committed
740
        super().__init__()
741
742
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
wangding zeng's avatar
wangding zeng committed
743
744
        self.config = config
        self.quant_config = quant_config
745
746
747
748
749
750
751
752
753
754
755
756
757

        # `packed_modules_mapping` needs to be modified before
        # initializing DeepseekV2Model, as it is passed inplace to
        # quantization config init and may be used to select the
        # quant_method for relevant layers during initialization.
        self.fuse_qkv_a_proj = hasattr(
            config, "q_lora_rank") and config.q_lora_rank is not None
        if self.fuse_qkv_a_proj:
            self.packed_modules_mapping["fused_qkv_a_proj"] = [
                "q_a_proj",
                "kv_a_proj_with_mqa",
            ]

758
759
        self.model = DeepseekV2Model(vllm_config=vllm_config,
                                     prefix=maybe_prefix(prefix, "model"))
760
761
762
763
764
765
        if get_pp_group().is_last_rank:
            self.lm_head = ParallelLMHead(config.vocab_size,
                                          config.hidden_size,
                                          quant_config=quant_config)
        else:
            self.lm_head = PPMissingLayer()
wangding zeng's avatar
wangding zeng committed
766
        self.logits_processor = LogitsProcessor(config.vocab_size)
767
768
        self.make_empty_intermediate_tensors = (
            self.model.make_empty_intermediate_tensors)
769
770
771
772
773
774
775
776
        self.expert_weights = []

        # Set MoE hyperparameters
        self.num_moe_layers = (config.num_hidden_layers -
                               config.first_k_dense_replace)
        self.num_expert_groups = config.n_group

        self.moe_layers: list[FusedMoE] = []
777
        example_moe = None
778
        for layer in self.model.layers:
779
780
781
            if isinstance(layer, PPMissingLayer):
                continue

782
783
            assert isinstance(layer, DeepseekV2DecoderLayer)
            if isinstance(layer.mlp, DeepseekV2MoE):
784
785
                # Pick last one layer since the first ones may be dense layers.
                example_moe = layer.mlp
786
787
                self.moe_layers.append(layer.mlp.experts)

788
789
790
        if example_moe is None:
            raise RuntimeError("No DeepseekV2MoE layer found in model.layers.")

791
792
793
794
795
796
        self.num_logical_experts = example_moe.n_logical_experts
        self.num_physical_experts = example_moe.n_physical_experts
        self.num_local_physical_experts = example_moe.n_local_physical_experts
        self.num_routed_experts = example_moe.n_routed_experts
        self.num_shared_experts = example_moe.n_shared_experts
        self.num_redundant_experts = example_moe.n_redundant_experts
797
798
799
800
801
        
        self.quant_method = None
        if quant_config is not None:
            self.quant_method=quant_config.get_name()
        self.use_llama_nn = os.environ.get('LLAMA_NN') == '1'
wangding zeng's avatar
wangding zeng committed
802

803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
    def set_eplb_state(
        self,
        expert_load_view: torch.Tensor,
        logical_to_physical_map: torch.Tensor,
        logical_replica_count: torch.Tensor,
    ) -> None:
        for layer_idx, layer in enumerate(self.moe_layers):
            # Register the expert weights.
            self.expert_weights.append(layer.get_expert_weights())
            layer.set_eplb_state(
                moe_layer_idx=layer_idx,
                expert_load_view=expert_load_view,
                logical_to_physical_map=logical_to_physical_map,
                logical_replica_count=logical_replica_count,
            )
818

819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
    def update_physical_experts_metadata(
        self,
        num_physical_experts: int,
        num_local_physical_experts: int,
    ) -> None:
        assert self.num_local_physical_experts == num_local_physical_experts
        self.num_physical_experts = num_physical_experts
        self.num_local_physical_experts = num_local_physical_experts
        self.num_redundant_experts = (num_physical_experts -
                                      self.num_logical_experts)
        for layer in self.model.layers:
            if isinstance(layer.mlp, DeepseekV2MoE):
                moe = layer.mlp
                moe.n_local_physical_experts = num_local_physical_experts
                moe.n_physical_experts = num_physical_experts
                moe.n_redundant_experts = self.num_redundant_experts
                moe.experts.update_expert_map()

837
838
839
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.model.get_input_embeddings(input_ids)

wangding zeng's avatar
wangding zeng committed
840
841
842
843
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
844
        intermediate_tensors: Optional[IntermediateTensors] = None,
845
        inputs_embeds: Optional[torch.Tensor] = None,
846
    ) -> Union[torch.Tensor, IntermediateTensors]:
847
        hidden_states = self.model(input_ids, positions, intermediate_tensors,
848
                                   inputs_embeds)
wangding zeng's avatar
wangding zeng committed
849
850
        return hidden_states

851
852
853
854
855
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
856
        logits = self.logits_processor(self.lm_head, hidden_states,
wangding zeng's avatar
wangding zeng committed
857
858
859
                                       sampling_metadata)
        return logits

860
861
    def load_weights(self, weights: Iterable[tuple[str,
                                                   torch.Tensor]]) -> set[str]:
wangding zeng's avatar
wangding zeng committed
862
863
864
865
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
            ("gate_up_proj", "gate_proj", 0),
            ("gate_up_proj", "up_proj", 1),
866
867
            ("fused_qkv_a_proj", "q_a_proj", 0),
            ("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1),
wangding zeng's avatar
wangding zeng committed
868
869
        ]

870
871
872
873
874
875
        # Params for weights, fp8 weight scales, fp8 activation scales
        # (param_name, weight_name, expert_id, shard_id)
        expert_params_mapping = FusedMoE.make_expert_params_mapping(
            ckpt_gate_proj_name="gate_proj",
            ckpt_down_proj_name="down_proj",
            ckpt_up_proj_name="up_proj",
876
877
            num_experts=self.config.n_routed_experts,
            num_redundant_experts=self.num_redundant_experts)
878

wangding zeng's avatar
wangding zeng committed
879
        params_dict = dict(self.named_parameters())
880
        loaded_params: set[str] = set()
wangding zeng's avatar
wangding zeng committed
881
882
883
        for name, loaded_weight in weights:
            if "rotary_emb.inv_freq" in name:
                continue
884

885
886
887
            spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
            if spec_layer is not None:
                continue  # skip spec decode layers for main model
888

wangding zeng's avatar
wangding zeng committed
889
            for (param_name, weight_name, shard_id) in stacked_params_mapping:
890
                # Skip non-stacked layers and experts (experts handled below).
wangding zeng's avatar
wangding zeng committed
891
892
                if weight_name not in name:
                    continue
893
894
895
896
897
898
899
900
                # We have mlp.experts[0].gate_proj in the checkpoint.
                # Since we handle the experts below in expert_params_mapping,
                # we need to skip here BEFORE we update the name, otherwise
                # name will be updated to mlp.experts[0].gate_up_proj, which
                # will then be updated below in expert_params_mapping
                # for mlp.experts[0].gate_gate_up_proj, which breaks load.
                if (("mlp.experts." in name) and name not in params_dict):
                    continue
901
                name_mapped = name.replace(weight_name, param_name)
902
903
904

                # QKV fusion is optional, fall back to normal
                # weight loading if it's not enabled
905
                # if go with fusion option, then update name
906
                if ((param_name == "fused_qkv_a_proj")
907
                        and name_mapped not in params_dict):
908
                    continue
909
910
                else:
                    name = name_mapped
wangding zeng's avatar
wangding zeng committed
911
912
913
                # Skip loading extra bias for GPTQ models.
                if name.endswith(".bias") and name not in params_dict:
                    continue
914
915
916
917

                if is_pp_missing_parameter(name, self):
                    continue

wangding zeng's avatar
wangding zeng committed
918
919
920
921
922
                param = params_dict[name]
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                break
            else:
923
                is_expert_weight = False
924
925
926
927
                for mapping in expert_params_mapping:
                    param_name, weight_name, expert_id, shard_id = mapping
                    if weight_name not in name:
                        continue
928

929
930
931
932
933
934
935
936
937
                    # Anyway, this is an expert weight and should not be
                    # attempted to load as other weights later
                    is_expert_weight = True

                    # Do not modify `name` since the loop may continue here
                    # Instead, create a new variable
                    name_mapped = name.replace(weight_name, param_name)

                    if is_pp_missing_parameter(name_mapped, self):
938
939
                        continue

940
941
942
943
944
945
946
947
948
949
950
951
952
                    param = params_dict[name_mapped]
                    # We should ask the weight loader to return success or not
                    # here since otherwise we may skip experts with other
                    # available replicas.
                    weight_loader = typing.cast(Callable[..., bool],
                                                param.weight_loader)
                    success = weight_loader(param,
                                            loaded_weight,
                                            name_mapped,
                                            shard_id=shard_id,
                                            expert_id=expert_id,
                                            return_success=True)
                    if success:
953
                        name = name_mapped
954
                        break
955
                else:
956
957
958
959
960
961
                    if is_expert_weight:
                        # We've checked that this is an expert weight
                        # However it's not mapped locally to this rank
                        # So we simply skip it
                        continue

962
963
964
965
                    # Skip loading extra bias for GPTQ models.
                    if name.endswith(".bias") and name not in params_dict:
                        continue

966
967
968
969
970
                    # Remapping the name of FP8 kv-scale.
                    name = maybe_remap_kv_scale_name(name, params_dict)
                    if name is None:
                        continue

971
972
973
                    if is_pp_missing_parameter(name, self):
                        continue

974
975
976
977
                    param = params_dict[name]
                    weight_loader = getattr(param, "weight_loader",
                                            default_weight_loader)
                    weight_loader(param, loaded_weight)
978
            loaded_params.add(name)
979
980
981
982
983
984
985
            
        if self.use_llama_nn and self.quant_method is None:
                lay_key_words = [
                    "self_attn.q_a_proj.weight",
                    "self_attn.kv_a_proj_with_mqa.weight",
                    "mlp.gate.weight",
                    "mlp.gate_up_proj.weight",
zhuwenwen's avatar
zhuwenwen committed
986
987
988
                    "mlp.down_proj.weight",
                    "shared_experts.gate_up_proj.weight",
                    "shared_experts.down_proj.weight",
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
                    "self_attn.q_proj.weight",
                    "self_attn.q_b_proj.weight",
                    "self_attn.kv_b_proj.weight",
                    "self_attn.o_proj.weight",
                    "lm_head.weight"
                ]

                combined_words = "|".join(lay_key_words)
                
                for layername in loaded_params:
                    weight = params_dict[layername]
                    matches = re.findall(combined_words, layername)
                    if matches:
                        _weight = torch.zeros_like(weight.data)
                        ori_shape =_weight.shape
                        
                        ops.trans_w16_gemm(_weight, weight.data, _weight.shape[0], _weight.shape[1])
                        weight.data.copy_(_weight)
                        
                        weight.data=weight.data.reshape(ori_shape[1],-1)
                        
1010
        return loaded_params
1011
1012
1013
1014


class DeepseekV3ForCausalLM(DeepseekV2ForCausalLM):
    pass
1015
1016


1017
1018
1019
1020
# Compatibility with
# https://huggingface.co/deepseek-ai/DeepSeek-V3-Base/blob/main/configuration_deepseek.py
def get_spec_layer_idx_from_weight_name(config: Union[DeepseekV2Config,
                                                      DeepseekV3Config],
1021
                                        weight_name: str) -> Optional[int]:
1022
1023
    if (hasattr(config, "num_nextn_predict_layers")
            and config.num_nextn_predict_layers > 0):
1024
1025
1026
1027
1028
        layer_idx = config.num_hidden_layers
        for i in range(config.num_nextn_predict_layers):
            if weight_name.startswith(f"model.layers.{layer_idx+i}."):
                return layer_idx + i
    return None