deepseek_v2.py 41.3 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0

wangding zeng's avatar
wangding zeng committed
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
# 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.
24
"""Inference-only DeepseekV2/DeepseekV3 model."""
王敏's avatar
王敏 committed
25
26
import os
import re
27
import vllm.envs as envs
28
from typing import Any, Dict, Iterable, Optional, Set, Tuple, Union
wangding zeng's avatar
wangding zeng committed
29
30
31
32
33

import torch
from torch import nn
from transformers import PretrainedConfig

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

58
59
from .interfaces import SupportsPP
from .utils import (PPMissingLayer, is_pp_missing_parameter,
60
61
                    make_empty_intermediate_tensors_factory, make_layers,
                    maybe_prefix)
王敏's avatar
王敏 committed
62
from vllm import _custom_ops as ops
63

64
65
66
67
from vllm.model_executor.layers.quantization.utils.int8_utils import (
    block_dequant as int8_block_dequant,
)

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

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

        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
121
        self.gate = ReplicatedLinear(config.hidden_size,
122
                                     config.n_routed_experts,
wangding zeng's avatar
wangding zeng committed
123
                                     bias=False,
124
125
                                     quant_config=None,
                                     prefix=f"{prefix}.gate")
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
        if config.topk_method == "noaux_tc":
            self.gate.e_score_correction_bias = nn.Parameter(
                torch.empty(config.n_routed_experts))
        else:
            self.gate.e_score_correction_bias = None

        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,
王敏's avatar
王敏 committed
145
            e_score_correction_bias=self.gate.e_score_correction_bias,)
146

wangding zeng's avatar
wangding zeng committed
147
148
149
150
151
152
153
154
155
        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,
                reduce_results=False,
156
                prefix=f"{prefix}.shared_experts",
wangding zeng's avatar
wangding zeng committed
157
            )
lizhigong's avatar
lizhigong committed
158
159
160
        from vllm.two_batch_overlap.two_batch_overlap import tbo_all_reduce, is_enable_tbo
        self.tbo_all_reduce = tbo_all_reduce
        self.enable_tbo = is_enable_tbo()
wangding zeng's avatar
wangding zeng committed
161
162
163
164

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        num_tokens, hidden_dim = hidden_states.shape
        hidden_states = hidden_states.view(-1, hidden_dim)
165
        if self.n_shared_experts is not None:
wangding zeng's avatar
wangding zeng committed
166
167
168
            shared_output = self.shared_experts(hidden_states)
        # router_logits: (num_tokens, n_experts)
        router_logits, _ = self.gate(hidden_states)
169
170
171
172
173
        # if hidden_states.dtype != torch.float16:
        #     final_hidden_states = self.experts(
        #         hidden_states=hidden_states,
        #         router_logits=router_logits) * self.routed_scaling_factor
        # else:
zhuwenwen's avatar
zhuwenwen committed
174
175
        #     # Fix FP16 overflow
        #     # See DeepseekV2DecoderLayer for more details.
176
177
        #     final_hidden_states = self.experts(hidden_states=hidden_states,
        #                                        router_logits=router_logits)
zhuwenwen's avatar
zhuwenwen committed
178
        
179
180
181
        final_hidden_states = self.experts(
            hidden_states=hidden_states,
            router_logits=router_logits) * self.routed_scaling_factor
182
        if shared_output is not None:
183
            final_hidden_states = final_hidden_states + shared_output
zhuwenwen's avatar
zhuwenwen committed
184
            
185
186
187
188
        # if shared_output is not None:
        #     if hidden_states.dtype != torch.float16:
        #         final_hidden_states = final_hidden_states + shared_output
        #     else:
zhuwenwen's avatar
zhuwenwen committed
189
190
        #         # Fix FP16 overflow
        #         # See DeepseekV2DecoderLayer for more details.
191
192
        #         final_hidden_states = final_hidden_states + shared_output \
        #             * (1. / self.routed_scaling_factor)
193
        if self.tp_size > 1:
lizhigong's avatar
lizhigong committed
194
195
196
197
198
            if self.enable_tbo:
                final_hidden_states = self.tbo_all_reduce(final_hidden_states)
            else:
                final_hidden_states = tensor_model_parallel_all_reduce(
                    final_hidden_states)
wangding zeng's avatar
wangding zeng committed
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226

        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,
        config: PretrainedConfig,
        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,
        rope_scaling: Optional[Dict[str, Any]] = None,
        max_position_embeddings: int = 8192,
        cache_config: Optional[CacheConfig] = None,
        quant_config: Optional[QuantizationConfig] = None,
227
        prefix: str = "",
wangding zeng's avatar
wangding zeng committed
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
    ) -> 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,
249
250
                                             quant_config=quant_config,
                                             prefix=f"{prefix}.q_a_proj")
wangding zeng's avatar
wangding zeng committed
251
252
253
254
255
256
            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,
257
258
                                                 quant_config=quant_config,
                                                 prefix=f"{prefix}.q_b_proj")
wangding zeng's avatar
wangding zeng committed
259
260
261
262
263
        else:
            self.q_proj = ColumnParallelLinear(self.hidden_size,
                                               self.num_heads *
                                               self.qk_head_dim,
                                               bias=False,
264
265
                                               quant_config=quant_config,
                                               prefix=f"{prefix}.q_proj")
wangding zeng's avatar
wangding zeng committed
266

267
268
269
270
271
272
        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
273
274
275
276
277
278
        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,
279
280
            quant_config=quant_config,
            prefix=f"{prefix}.kv_b_proj")
wangding zeng's avatar
wangding zeng committed
281
282
283
284
        # O projection.
        self.o_proj = RowParallelLinear(self.num_heads * self.v_head_dim,
                                        self.hidden_size,
                                        bias=False,
285
286
                                        quant_config=quant_config,
                                        prefix=f"{prefix}.o_proj")
287
288
        if rope_scaling:
            rope_scaling["rope_type"] = 'deepseek_yarn'
289

wangding zeng's avatar
wangding zeng committed
290
291
292
293
294
295
296
297
298
299
300
301
302
303
        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,
304
                              self.qk_head_dim,
wangding zeng's avatar
wangding zeng committed
305
306
307
                              self.scaling,
                              num_kv_heads=self.num_local_heads,
                              cache_config=cache_config,
308
309
                              quant_config=quant_config,
                              prefix=f"{prefix}.attn")
wangding zeng's avatar
wangding zeng committed
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335

    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)
        kv_a = self.kv_a_layernorm(kv_a.contiguous())
        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:]
336

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

wangding zeng's avatar
wangding zeng committed
339
340
341
342
        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
343
344
345
346
        # 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)
347
        attn_output = self.attn(q, k, v)
wangding zeng's avatar
wangding zeng committed
348
        attn_output = attn_output.view(
349
350
            -1, self.num_local_heads,
            self.qk_head_dim)[..., :self.v_head_dim].reshape(
wangding zeng's avatar
wangding zeng committed
351
352
353
354
355
                -1, self.num_local_heads * self.v_head_dim)
        output, _ = self.o_proj(attn_output)
        return output


356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
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
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
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,
        config: PretrainedConfig,
        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,
        rope_scaling: Optional[Dict[str, Any]] = None,
        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:
            self.q_a_proj = ReplicatedLinear(self.hidden_size,
                                             self.q_lora_rank,
                                             bias=False,
                                             quant_config=quant_config,
                                             prefix=f"{prefix}.q_a_proj")
            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,
                                                 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_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")
        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")

442
443
        if rope_scaling:
            rope_scaling["rope_type"] = 'deepseek_yarn'
444
445
446
447
448
449
450
451
452
453
454
455
        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

456
457
458
459
460
461
        # 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
462
463
        self.mla_attn = Attention(
            num_heads=self.num_local_heads,
464
            head_size=self.kv_lora_rank + self.qk_rope_head_dim,
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
            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,
            rotary_emb=self.rotary_emb,
            q_proj=self.q_proj if self.q_lora_rank is None else self.q_b_proj,
            kv_b_proj=self.kv_b_proj,
            o_proj=self.o_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:
        if self.q_lora_rank is not None:
            ckq = self.q_a_proj(hidden_states)[0]
            hidden_states_or_q_c = self.q_a_layernorm(ckq)
        else:
            hidden_states_or_q_c = hidden_states
        kv_c, k_pe = self.kv_a_proj_with_mqa(hidden_states)[0].split(
            [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
        kv_c_normed = self.kv_a_layernorm(kv_c.contiguous())
500
501
502
503
        return self.mla_attn(hidden_states_or_q_c,
                             kv_c_normed,
                             k_pe,
                             output_shape=hidden_states.shape)
504
505


wangding zeng's avatar
wangding zeng committed
506
507
508
509
510
class DeepseekV2DecoderLayer(nn.Module):

    def __init__(
        self,
        config: PretrainedConfig,
511
        prefix: str,
512
        model_config: ModelConfig,
wangding zeng's avatar
wangding zeng committed
513
514
515
516
517
518
519
520
521
        cache_config: Optional[CacheConfig] = None,
        quant_config: Optional[QuantizationConfig] = None,
    ) -> 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)
522
523
524
        # DecoderLayers are created with `make_layers` which passes the prefix
        # with the layer's index.
        layer_idx = int(prefix.split(sep='.')[-1])
525
        self.layer_idx = layer_idx
526
527
528
529
530
        if model_config.use_mla:
            attn_cls = DeepseekV2MLAAttention
        else:
            attn_cls = DeepseekV2Attention
        self.self_attn = attn_cls(
wangding zeng's avatar
wangding zeng committed
531
532
533
534
535
536
537
538
539
540
541
542
543
544
            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,
545
            prefix=f"{prefix}.self_attn",
wangding zeng's avatar
wangding zeng committed
546
        )
547

wangding zeng's avatar
wangding zeng committed
548
549
550
        if (config.n_routed_experts is not None
                and layer_idx >= config.first_k_dense_replace
                and layer_idx % config.moe_layer_freq == 0):
551
552
553
554
555
            self.mlp = DeepseekV2MoE(
                config=config,
                quant_config=quant_config,
                prefix=f"{prefix}.mlp",
            )
wangding zeng's avatar
wangding zeng committed
556
557
558
559
560
561
        else:
            self.mlp = DeepseekV2MLP(
                hidden_size=config.hidden_size,
                intermediate_size=config.intermediate_size,
                hidden_act=config.hidden_act,
                quant_config=quant_config,
562
                prefix=f"{prefix}.mlp",
wangding zeng's avatar
wangding zeng committed
563
564
565
566
567
            )
        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)
568
        self.routed_scaling_factor = config.routed_scaling_factor
wangding zeng's avatar
wangding zeng committed
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587

    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,
        )

zhuwenwen's avatar
zhuwenwen committed
588
589
590
591
592
593
594
595
596
        # 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.
        #     hidden_states *= 1. / self.routed_scaling_factor
        #     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
597

wangding zeng's avatar
wangding zeng committed
598
599
600
601
        # Fully Connected
        hidden_states, residual = self.post_attention_layernorm(
            hidden_states, residual)
        hidden_states = self.mlp(hidden_states)
602

zhuwenwen's avatar
zhuwenwen committed
603
604
605
606
607
608
609
        # 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
610
        #     hidden_states *= 1. / self.routed_scaling_factor
611

wangding zeng's avatar
wangding zeng committed
612
613
614
        return hidden_states, residual


615
@support_torch_compile
wangding zeng's avatar
wangding zeng committed
616
617
618
619
class DeepseekV2Model(nn.Module):

    fall_back_to_pt_during_load = False

620
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
wangding zeng's avatar
wangding zeng committed
621
        super().__init__()
622
623

        config = vllm_config.model_config.hf_config
624
        model_config = vllm_config.model_config
625
626
        cache_config = vllm_config.cache_config
        quant_config = vllm_config.quant_config
627
        self.config = config
628

wangding zeng's avatar
wangding zeng committed
629
630
        self.vocab_size = config.vocab_size

631
632
633
634
        if get_pp_group().is_first_rank:
            self.embed_tokens = VocabParallelEmbedding(
                config.vocab_size,
                config.hidden_size,
635
636
                quant_config=quant_config,
                prefix=f"{prefix}.embed_tokens")
637
638
639
640
641
642
643
644
        else:
            self.embed_tokens = PPMissingLayer()

        self.start_layer, self.end_layer, self.layers = make_layers(
            config.num_hidden_layers,
            lambda prefix: DeepseekV2DecoderLayer(
                config,
                prefix,
645
                model_config=model_config,
646
647
648
649
650
651
652
653
654
                cache_config=cache_config,
                quant_config=quant_config,
            ),
            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()
655
656
657
        self.make_empty_intermediate_tensors = (
            make_empty_intermediate_tensors_factory(
                ["hidden_states", "residual"], config.hidden_size))
wangding zeng's avatar
wangding zeng committed
658

659
660
661
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embed_tokens(input_ids)

wangding zeng's avatar
wangding zeng committed
662
663
664
665
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
666
        intermediate_tensors: Optional[IntermediateTensors],
667
        inputs_embeds: Optional[torch.Tensor] = None,
668
    ) -> Union[torch.Tensor, IntermediateTensors]:
669
        if get_pp_group().is_first_rank:
670
671
672
673
            if inputs_embeds is not None:
                hidden_states = inputs_embeds
            else:
                hidden_states = self.get_input_embeddings(input_ids)
674
675
676
677
678
679
            residual = None
        else:
            assert intermediate_tensors is not None
            hidden_states = intermediate_tensors["hidden_states"]
            residual = intermediate_tensors["residual"]

680
        for layer in self.layers[self.start_layer:self.end_layer]:
681
            hidden_states, residual = layer(positions, hidden_states, residual)
682
683
684
685
686
687
688

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

wangding zeng's avatar
wangding zeng committed
689
690
691
692
        hidden_states, _ = self.norm(hidden_states, residual)
        return hidden_states


693
class DeepseekV2ForCausalLM(nn.Module, SupportsPP):
wangding zeng's avatar
wangding zeng committed
694

695
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
wangding zeng's avatar
wangding zeng committed
696
        super().__init__()
697
698
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
699
700
701
702
703
704
705

        self.quant_method = None
        if quant_config is not None:
            self.quant_method = quant_config.get_name()
            os.environ['LLAMA_NN'] = '0'
            os.environ['LM_NN'] = '0'

706
        self.use_w4a16_moe_sz = os.environ.get('AWQ_MOE_SZ') == '1'
zhuwenwen's avatar
zhuwenwen committed
707
        self.tritonsingleton= W8a8GetCacheJSON()  
wangding zeng's avatar
wangding zeng committed
708
709
        self.config = config
        self.quant_config = quant_config
王敏's avatar
王敏 committed
710

711
        self.model = DeepseekV2Model(vllm_config=vllm_config,
712
                                     prefix=maybe_prefix(prefix, "model"))
713
714
715
716
717
718
        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
719
        self.logits_processor = LogitsProcessor(config.vocab_size)
720
721
        self.make_empty_intermediate_tensors = (
            self.model.make_empty_intermediate_tensors)
王敏's avatar
王敏 committed
722
        self.use_llama_nn = os.environ.get('LLAMA_NN') == '1'
723
724
        self.use_awq_pad = os.environ.get('AWQ_PAD') == '1'
        
725
726
727
    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
728
729
730
731
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
732
        intermediate_tensors: Optional[IntermediateTensors] = None,
733
        inputs_embeds: Optional[torch.Tensor] = None,
734
    ) -> Union[torch.Tensor, IntermediateTensors]:
735
        hidden_states = self.model(input_ids, positions, intermediate_tensors,
736
                                   inputs_embeds)
wangding zeng's avatar
wangding zeng committed
737
738
        return hidden_states

739
740
741
742
743
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
744
        logits = self.logits_processor(self.lm_head, hidden_states,
wangding zeng's avatar
wangding zeng committed
745
746
747
                                       sampling_metadata)
        return logits

748
749
750
751
752
753
754
755
756
757
758
759
760
    def make_empty_intermediate_tensors(
            self, batch_size: int, dtype: torch.dtype,
            device: torch.device) -> IntermediateTensors:
        return IntermediateTensors({
            "hidden_states":
            torch.zeros((batch_size, self.config.hidden_size),
                        dtype=dtype,
                        device=device),
            "residual":
            torch.zeros((batch_size, self.config.hidden_size),
                        dtype=dtype,
                        device=device),
        })
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
        
    def restore_qzeros_tensor(self, qzeros, qscales):

        low_bits = qzeros & 0x0F
        high_bits = qzeros >> 4
        
        zeors_tensor = torch.stack([low_bits, high_bits], dim=2).view(qzeros.shape[0], -1 , qzeros.shape[-1])
        zeors_int16 = zeors_tensor.to(torch.int16)
        assert zeors_int16.shape == qscales.shape

        uint16_tensor1 = zeors_int16.view(torch.uint16)
        uint16_tensor2 = qscales.view(torch.uint16)
        
        uint32_tensor1 = uint16_tensor1.to(torch.int32) << 16
        uint32_tensor2 = uint16_tensor2.to(torch.int32)
        
        result_tensor = uint32_tensor1 + uint32_tensor2
        result_tensor =result_tensor.view(torch.uint32)
        result_tensor = result_tensor.transpose(1, 2).contiguous()
        return result_tensor
781

782
783
    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
wangding zeng's avatar
wangding zeng committed
784
785
786
787
788
789
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
            ("gate_up_proj", "gate_proj", 0),
            ("gate_up_proj", "up_proj", 1),
        ]

790
791
        # Params for weights, fp8 weight scales, fp8 activation scales
        # (param_name, weight_name, expert_id, shard_id)
王敏's avatar
王敏 committed
792
793
794
795
796
        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",
            num_experts=self.config.n_routed_experts)
797

wangding zeng's avatar
wangding zeng committed
798
        params_dict = dict(self.named_parameters())
799
        loaded_params: Set[str] = set()
wangding zeng's avatar
wangding zeng committed
800
801
802
        for name, loaded_weight in weights:
            if "rotary_emb.inv_freq" in name:
                continue
803

804
805
806
            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
807

wangding zeng's avatar
wangding zeng committed
808
            for (param_name, weight_name, shard_id) in stacked_params_mapping:
809
                # Skip non-stacked layers and experts (experts handled below).
wangding zeng's avatar
wangding zeng committed
810
811
                if weight_name not in name:
                    continue
812
813
814
815
816
817
818
819
                # 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
wangding zeng's avatar
wangding zeng committed
820
821
822
823
                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
824
825
826
827

                if is_pp_missing_parameter(name, self):
                    continue

wangding zeng's avatar
wangding zeng committed
828
829
830
831
832
                param = params_dict[name]
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                break
            else:
833
834
835
836
837
                for mapping in expert_params_mapping:
                    param_name, weight_name, expert_id, shard_id = mapping
                    if weight_name not in name:
                        continue
                    name = name.replace(weight_name, param_name)
838
839
840
841

                    if is_pp_missing_parameter(name, self):
                        continue

842
843
844
845
                    param = params_dict[name]
                    weight_loader = param.weight_loader
                    weight_loader(param,
                                  loaded_weight,
846
                                  name,
847
848
849
850
851
852
853
854
                                  shard_id=shard_id,
                                  expert_id=expert_id)
                    break
                else:
                    # Skip loading extra bias for GPTQ models.
                    if name.endswith(".bias") and name not in params_dict:
                        continue

855
856
857
858
859
                    # Remapping the name of FP8 kv-scale.
                    name = maybe_remap_kv_scale_name(name, params_dict)
                    if name is None:
                        continue

860
861
862
                    if is_pp_missing_parameter(name, self):
                        continue

863
864
865
866
                    param = params_dict[name]
                    weight_loader = getattr(param, "weight_loader",
                                            default_weight_loader)
                    weight_loader(param, loaded_weight)
867
            loaded_params.add(name)
王敏's avatar
王敏 committed
868
            
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
        if self.use_llama_nn and self.quant_method is None:
            lay_key_words = [
                "self_attn.q_proj.weight",
                "self_attn.q_a_proj.weight",
                "self_attn.q_b_proj.weight",
                "self_attn.kv_a_proj_with_mqa.weight",
                "self_attn.kv_b_proj.weight",
                "self_attn.o_proj.weight",
                "mlp.gate_up_proj.weight",
                "mlp.down_proj.weight",
                "mlp.gate.weight",
                "shared_experts.gate_up_proj.weight",
                "shared_experts.down_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)
zhuwenwen's avatar
zhuwenwen committed
898

zhuwenwen's avatar
zhuwenwen committed
899
        if hasattr(self.config, "quantization_config") and self.config.quantization_config["quant_method"] == "awq" and not envs.VLLM_USE_TRITON_AWQ:
900
901
902
903
            lay_key_words = [
                "self_attn.q_a_proj.qweight",
                "self_attn.q_b_proj.qweight",
                "self_attn.kv_b_proj.qweight",
zhuwenwen's avatar
zhuwenwen committed
904
                "self_attn.kv_a_proj_with_mqa.qweight",
905
906
907
908
909
910
911
                "self_attn.o_proj.qweight",
                "mlp.gate_up_proj.qweight",
                "mlp.down_proj.qweight",
                "mlp.shared_experts.gate_up_proj.qweight",
                "mlp.shared_experts.down_proj.qweight"
            ]
            combined_words = "|".join(lay_key_words)
912
913
914
            # moe_gather_sz
            moe_key_words = ["mlp.experts.w13_qweight", "mlp.experts.w2_qweight"]
            moe_combined_words = "|".join(moe_key_words)
915
916
917
918
919
            
            for layername in loaded_params:
                weight = params_dict[layername]
                
                matches = re.findall(combined_words, layername)
920

921
922
923
924
925
926
927
928
929
                        
                if self.use_w4a16_moe_sz:
                    matches_moe = re.findall(moe_combined_words, layername)
                    # sz.shape == s.shape.T
                    if matches_moe:
                        qzeros=params_dict[layername.replace("qweight", "qzeros")]
                        scales=params_dict[layername.replace("qweight", "scales")]
                        sz_tensor = self.restore_qzeros_tensor(qzeros, scales)
                        scales.data = sz_tensor
zhuwenwen's avatar
zhuwenwen committed
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
            
        if hasattr(self.config, "quantization_config") and self.config.quantization_config["quant_method"] == "blockwise_int8":
            lay_key_words = [
                "self_attn.q_a_proj.weight",
                "self_attn.q_b_proj.weight",
                "self_attn.kv_b_proj.weight",
                "self_attn.kv_a_proj_with_mqa.weight",
                "self_attn.o_proj.weight",
                "mlp.gate_up_proj.weight",
                "mlp.down_proj.weight",
                "mlp.shared_experts.gate_up_proj.weight",
                "mlp.shared_experts.down_proj.weight"
            ]
            combined_words = "|".join(lay_key_words)
            weight_shapes=[]
            all_json={}
            matched_key_words=set()
            
            for layername, weight in params_dict.items():  
                matches = re.findall(combined_words, layername)
                if matches and "scale" not in layername:
                    weight_data =params_dict[layername]
                    n=weight_data.shape[0]
                    
                    if len(matched_key_words) < 9 and matches[0] not in matched_key_words:
                        matched_key_words.add(matches[0])
                        k=weight_data.shape[1]
                        weight_shapes.append({n,k})
                        #print("n:{},k:{}".format(n,k))
                
                        json_file=self.tritonsingleton.get_blockint8json_name(n,k,128,128)
                        configs_dict=self.tritonsingleton.get_blockint8_triton_cache(json_file,n,k,128,128)
                        if configs_dict:
                            all_json.update(configs_dict)
                            
            self.tritonsingleton.triton_json_dict.append(all_json) 
            #print("self.tritonsingleton.triton_json_dict[0].shape:",len(self.tritonsingleton.triton_json_dict[0]))    
            for key, value in all_json.items():
                m=int(key.split('_')[0])
                n=int(key.split('_')[1])
                k=int(key.split('_')[2])
                # ops.triton_int8_gemm_helper(m=m,n=n,k=k,per_token_act_quant=True,per_out_channel_weight_quant=True,use_bias=False,best_config=value)
王敏's avatar
王敏 committed
972

973
        return loaded_params
974
975
976
977


class DeepseekV3ForCausalLM(DeepseekV2ForCausalLM):
    pass
978
979
980
981
982
983
984
985
986
987
988
989


def get_spec_layer_idx_from_weight_name(config: PretrainedConfig,
                                        weight_name: str) -> Optional[int]:
    if hasattr(config,
               "num_nextn_predict_layers") and (config.num_nextn_predict_layers
                                                > 0):
        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