deepseek_v2.py 54.9 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."""
王敏's avatar
王敏 committed
26
27
import os
import re
28
import vllm.envs as envs
zhuwenwen's avatar
zhuwenwen committed
29

30
31
import typing
from collections.abc import Callable, Iterable
32
from typing import Any, Optional, Union, Tuple
wangding zeng's avatar
wangding zeng committed
33
34
35
36
37

import torch
from torch import nn
from transformers import PretrainedConfig

38
from vllm.attention import Attention
39
from vllm.compilation.decorators import support_torch_compile
40
41
from vllm.config import (CacheConfig, ModelConfig, VllmConfig,
                         get_current_vllm_config)
王敏's avatar
王敏 committed
42
from vllm.distributed import (get_ep_group, get_pp_group, get_dp_group,
43
                              get_tensor_model_parallel_world_size)
wangding zeng's avatar
wangding zeng committed
44
from vllm.model_executor.layers.activation import SiluAndMul
45
from vllm.model_executor.layers.fused_moe import FusedMoE
wangding zeng's avatar
wangding zeng committed
46
47
48
49
50
51
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
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, SupportsPP
62
from .utils import (PPMissingLayer, is_pp_missing_parameter,
63
64
                    make_empty_intermediate_tensors_factory, make_layers,
                    maybe_prefix)
王敏's avatar
王敏 committed
65
from vllm import _custom_ops as ops
66
from vllm.utils import W8a8GetCacheJSON
67

wangding zeng's avatar
wangding zeng committed
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
        if hidden_act != "silu":
            raise ValueError(f"Unsupported activation: {hidden_act}. "
                             "Only silu is supported for now.")
        self.act_fn = SiluAndMul()

96
97
98
    def forward(self, x,
                rms_weight: Optional[torch.Tensor] = None,
                residual: Optional[torch.Tensor] = None,
99
100
                update_hd: Optional[bool] = False,
                xqxs: Optional[tuple[torch.Tensor, torch.Tensor]] = None):
101
102
        if envs.USE_FUSED_RMS_QUANT:
            gate_up, new_resi, _  = self.gate_up_proj(x, rms_weight, residual, update_hd=update_hd)
103
104
105
106
107
108
            if envs.USE_FUSED_SILU_MUL_QUANT:
                x, _ = self.down_proj(gate_up, use_fused_silu_mul_quant=True)
            else:
                x = self.act_fn(gate_up)
                x, _ = self.down_proj(x)
                
109
            return x, new_resi
110
111
112
113
114
        elif envs.USE_FUSED_CUSTOM_ALL_REDUCE_RMS_QUANT and xqxs is not None:
            gate_up, _ = self.gate_up_proj(x, xqxs=xqxs)
            x = self.act_fn(gate_up)
            x, _ = self.down_proj(x)
            return x
115
116
117
118
119
        else:
            gate_up, _ = self.gate_up_proj(x)
            x = self.act_fn(gate_up)
            x, _ = self.down_proj(x)
            return x
wangding zeng's avatar
wangding zeng committed
120
121
122
123
124
125
126
127


class DeepseekV2MoE(nn.Module):

    def __init__(
        self,
        config: PretrainedConfig,
        quant_config: Optional[QuantizationConfig] = None,
128
        prefix: str = "",
129
        enable_eplb: bool = False,
wangding zeng's avatar
wangding zeng committed
130
131
132
133
    ):
        super().__init__()
        self.tp_size = get_tensor_model_parallel_world_size()
        self.routed_scaling_factor = config.routed_scaling_factor
134
135
136
137
138
139

        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
140
141
142
143
144

        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
145
        self.gate = ReplicatedLinear(config.hidden_size,
146
                                     config.n_routed_experts,
wangding zeng's avatar
wangding zeng committed
147
                                     bias=False,
148
149
                                     quant_config=None,
                                     prefix=f"{prefix}.gate")
150
151
152
153
154
155
        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

156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
        # Load balancing settings.
        vllm_config = get_current_vllm_config()
        parallel_config = vllm_config.parallel_config
        self.enable_eplb = enable_eplb

        self.n_redundant_experts = parallel_config.num_redundant_experts
        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)
王敏's avatar
王敏 committed
171
        
王敏's avatar
王敏 committed
172
        self.experts = FusedMoE(
173
174
175
176
177
178
179
180
181
182
183
184
            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
185
            e_score_correction_bias=self.gate.e_score_correction_bias,
186
            enable_eplb=self.enable_eplb,
zhuwenwen's avatar
zhuwenwen committed
187
            num_redundant_experts=self.n_redundant_experts,
王敏's avatar
王敏 committed
188
            routed_scaling_factor=self.routed_scaling_factor)
189

wangding zeng's avatar
wangding zeng committed
190
191
192
        if config.n_shared_experts is not None:
            intermediate_size = (config.moe_intermediate_size *
                                 config.n_shared_experts)
王敏's avatar
王敏 committed
193
            self.shared_experts = DeepseekV2MLP(
wangding zeng's avatar
wangding zeng committed
194
195
196
197
                hidden_size=config.hidden_size,
                intermediate_size=intermediate_size,
                hidden_act=config.hidden_act,
                quant_config=quant_config,
198
199
                reduce_results=self.experts.must_reduce_shared_expert_outputs(
                ),
200
                prefix=f"{prefix}.shared_experts",
wangding zeng's avatar
wangding zeng committed
201
            )
王敏's avatar
王敏 committed
202

203
        from vllm.two_batch_overlap.two_batch_overlap import tbo_all_reduce
lizhigong's avatar
lizhigong committed
204
        self.tbo_all_reduce = tbo_all_reduce
wangding zeng's avatar
wangding zeng committed
205

206
207
    def forward(self, hidden_states: torch.Tensor,
                rms_weight: Optional[torch.Tensor] = None,
208
209
                residual: Optional[torch.Tensor] = None,
                xqxs: Optional[tuple[torch.Tensor, torch.Tensor]] = None
210
                ) -> torch.Tensor:
211
212
213
214
215
216
217
218
219
220
221
222
223
        if envs.USE_FUSED_CUSTOM_ALL_REDUCE_RMS_QUANT and xqxs is not None:
            num_tokens, hidden_dim = hidden_states.shape
            hidden_states = hidden_states.view(-1, hidden_dim)
            if self.n_shared_experts is not None:
                shared_output = self.shared_experts(hidden_states, xqxs=xqxs)
                    
            router_logits, _ = self.gate(hidden_states)

            if envs.VLLM_USE_LIGHTOP_MOE_SUM_MUL_ADD:
                final_hidden_states = self.experts(
                    hidden_states=hidden_states,
                    router_logits=router_logits,
                    shared_output=shared_output)
王敏's avatar
王敏 committed
224
            else:
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
                if hidden_states.dtype != torch.float16:
                    final_hidden_states = self.experts(
                        hidden_states=hidden_states,
                        router_logits=router_logits) * self.routed_scaling_factor
                else:
                    # Fix FP16 overflow
                    # See DeepseekV2DecoderLayer for more details.
                    final_hidden_states = self.experts(hidden_states=hidden_states,
                                                    router_logits=router_logits)
                
                if shared_output is not None:
                    if hidden_states.dtype != torch.float16:
                        final_hidden_states = final_hidden_states + shared_output
                    else:
                        # Fix FP16 overflow
                        # See DeepseekV2DecoderLayer for more details.
                        final_hidden_states = final_hidden_states + shared_output \
                            * (1. / self.routed_scaling_factor)

            if self.tp_size > 1:
                if envs.VLLM_ENABLE_TBO:
                    final_hidden_states = self.tbo_all_reduce(final_hidden_states)
                else:
                    final_hidden_states = (
                        self.experts.maybe_all_reduce_tensor_model_parallel(
                            final_hidden_states))
            return final_hidden_states.view(num_tokens, hidden_dim)
        else:        
            num_tokens, hidden_dim = hidden_states.shape
            hidden_states = hidden_states.view(-1, hidden_dim)
            if self.n_shared_experts is not None:
                if envs.USE_FUSED_RMS_QUANT:
                    shared_output, new_resi = self.shared_experts(hidden_states, rms_weight, residual, update_hd=True)
                else:
                    shared_output = self.shared_experts(hidden_states)
                    
            router_logits, _ = self.gate(hidden_states)
            
263

264
            if envs.VLLM_USE_LIGHTOP_MOE_SUM_MUL_ADD:
265
266
                final_hidden_states = self.experts(
                    hidden_states=hidden_states,
267
268
                    router_logits=router_logits,
                    shared_output=shared_output)
269
            else:
270
                if hidden_states.dtype != torch.float16:
271
272
273
                    final_hidden_states = self.experts(
                        hidden_states=hidden_states,
                        router_logits=router_logits) * self.routed_scaling_factor
274
275
276
                else:
                    # Fix FP16 overflow
                    # See DeepseekV2DecoderLayer for more details.
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
                    final_hidden_states = self.experts(hidden_states=hidden_states,
                                                    router_logits=router_logits)
            
                if shared_output is not None:
                    if hidden_states.dtype != torch.float16:
                        final_hidden_states = final_hidden_states + shared_output
                    else:
                        # Fix FP16 overflow
                        # See DeepseekV2DecoderLayer for more details.
                        final_hidden_states = final_hidden_states + shared_output \
                            * (1. / self.routed_scaling_factor)

            if self.tp_size > 1:
                if envs.VLLM_ENABLE_TBO:
                    final_hidden_states = self.tbo_all_reduce(final_hidden_states)
                else:
                    final_hidden_states = (
                        self.experts.maybe_all_reduce_tensor_model_parallel(
                            final_hidden_states))
296

297
298
            if envs.USE_FUSED_RMS_QUANT:
                return final_hidden_states.view(num_tokens, hidden_dim), new_resi
lizhigong's avatar
lizhigong committed
299
            else:
300
                return final_hidden_states.view(num_tokens, hidden_dim)
wangding zeng's avatar
wangding zeng committed
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322


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,
323
        rope_scaling: Optional[dict[str, Any]] = None,
wangding zeng's avatar
wangding zeng committed
324
325
326
        max_position_embeddings: int = 8192,
        cache_config: Optional[CacheConfig] = None,
        quant_config: Optional[QuantizationConfig] = None,
327
        prefix: str = "",
wangding zeng's avatar
wangding zeng committed
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
    ) -> 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,
349
350
                                             quant_config=quant_config,
                                             prefix=f"{prefix}.q_a_proj")
wangding zeng's avatar
wangding zeng committed
351
352
353
354
355
356
            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,
357
358
                                                 quant_config=quant_config,
                                                 prefix=f"{prefix}.q_b_proj")
wangding zeng's avatar
wangding zeng committed
359
360
361
362
363
        else:
            self.q_proj = ColumnParallelLinear(self.hidden_size,
                                               self.num_heads *
                                               self.qk_head_dim,
                                               bias=False,
364
365
                                               quant_config=quant_config,
                                               prefix=f"{prefix}.q_proj")
wangding zeng's avatar
wangding zeng committed
366

367
368
369
370
371
372
        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
373
374
375
376
377
378
        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,
379
380
            quant_config=quant_config,
            prefix=f"{prefix}.kv_b_proj")
wangding zeng's avatar
wangding zeng committed
381
382
383
384
        # O projection.
        self.o_proj = RowParallelLinear(self.num_heads * self.v_head_dim,
                                        self.hidden_size,
                                        bias=False,
385
386
                                        quant_config=quant_config,
                                        prefix=f"{prefix}.o_proj")
387
388
        if rope_scaling:
            rope_scaling["rope_type"] = 'deepseek_yarn'
389

wangding zeng's avatar
wangding zeng committed
390
391
392
393
394
395
396
397
398
399
400
401
402
403
        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,
404
                              self.qk_head_dim,
wangding zeng's avatar
wangding zeng committed
405
406
407
                              self.scaling,
                              num_kv_heads=self.num_local_heads,
                              cache_config=cache_config,
408
409
                              quant_config=quant_config,
                              prefix=f"{prefix}.attn")
wangding zeng's avatar
wangding zeng committed
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

    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:]
436

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

wangding zeng's avatar
wangding zeng committed
439
440
441
442
        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
443
444
445
446
        # 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)
447
        attn_output = self.attn(q, k, v)
wangding zeng's avatar
wangding zeng committed
448
        attn_output = attn_output.view(
449
450
            -1, self.num_local_heads,
            self.qk_head_dim)[..., :self.v_head_dim].reshape(
wangding zeng's avatar
wangding zeng committed
451
452
453
454
455
                -1, self.num_local_heads * self.v_head_dim)
        output, _ = self.o_proj(attn_output)
        return output


456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
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,
475
        rope_scaling: Optional[dict[str, Any]] = None,
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
        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:
501
502
            if envs.USE_FUSED_RMS_QUANT:
                self.q_a_proj = ReplicatedLinear(self.hidden_size,
503
504
505
                                             self.q_lora_rank,
                                             bias=False,
                                             quant_config=quant_config,
506
                                             eps=config.rms_norm_eps,
507
                                             prefix=f"{prefix}.q_a_proj")
508
                self.q_b_proj = ColumnParallelLinear(q_lora_rank,
509
510
511
512
                                                 self.num_heads *
                                                 self.qk_head_dim,
                                                 bias=False,
                                                 quant_config=quant_config,
513
                                                 eps=config.rms_norm_eps,
514
                                                 prefix=f"{prefix}.q_b_proj")
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
            else:
                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_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")
                
            self.q_a_layernorm = RMSNorm(self.q_lora_rank,
                                         eps=config.rms_norm_eps)

531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
        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")

559
560
        if rope_scaling:
            rope_scaling["rope_type"] = 'deepseek_yarn'
561
562
563
564
565
566
567
568
569
570
571
572
        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

573
574
575
576
577
578
        # 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
579
580
        self.mla_attn = Attention(
            num_heads=self.num_local_heads,
581
            head_size=self.kv_lora_rank + self.qk_rope_head_dim,
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
            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,
605
        rms_weight: Optional[torch.Tensor] = None,
606
607
608
609
610
611
612
613
614
615
        residual: Optional[torch.Tensor] = None,
        pa_rms_weight: Optional[torch.Tensor] = None,
        pa_residual: Optional[torch.Tensor] = None,
        pa_rms_eps: Optional[float] = 1e-6,
        pa_quant_dtype: Optional[torch.dtype] = torch.int8,
        update_input: Optional[bool] = True
    ) -> Union[torch.Tensor,
               Tuple[torch.Tensor, torch.Tensor],
               Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], 
               ]:
616
617
618
619
620
621
622
623
624
625
        if envs.USE_FUSED_RMS_QUANT and rms_weight is not None:
            if self.q_lora_rank is not None:
                q_c, new_residual, _, input_quant_args = self.q_a_proj(hidden_states, rms_weight=rms_weight, residual=residual, update_hd=False)
                q, _, _ = self.q_b_proj(q_c, rms_weight=self.q_a_layernorm.weight.data, residual=None, update_hd=False)
                
            else:
                q = self.q_proj(hidden_states)[0]
            kv_c, k_pe = self.kv_a_proj_with_mqa(hidden_states, quant_args=input_quant_args, update_hd=False)[0].split(
                                [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
            
626
627
628
629
            if envs.VLLM_USE_LIGHTOP:
                kv_c_normed = self.kv_a_layernorm.forward_cuda_opt(kv_c)
            else:
                kv_c_normed = self.kv_a_layernorm(kv_c.contiguous())
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644

            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)

            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], new_residual
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
        elif envs.USE_FUSED_CUSTOM_ALL_REDUCE_RMS_QUANT and pa_rms_weight is not None and pa_residual is not None:
            if self.q_lora_rank is not None:
                q_c = self.q_a_proj(hidden_states)[0]
                q_c = self.q_a_layernorm(q_c)
                q = self.q_b_proj(q_c)[0]
            else:
                q = self.q_proj(hidden_states)[0]
            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())

            q = q.view(-1, self.num_local_heads, self.qk_head_dim)
            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)

            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))
            packages_ = self.o_proj(attn_out, 
                                   pa_rms_weight=pa_rms_weight,
                                   pa_residual=pa_residual,
                                   pa_rms_eps=pa_rms_eps,
                                   pa_quant_dtype=pa_quant_dtype,
                                   update_input=update_input)[:4]
            assert len(packages_) == 4
            hs, resi, xq, xs = packages_
            assert xq is not None and xs is not None
            return hs, resi, xq, xs

679
        else:
680
681
682
683
684
685
686
687
            if self.q_lora_rank is not None:
                q_c = self.q_a_proj(hidden_states)[0]
                q_c = self.q_a_layernorm(q_c)
                q = self.q_b_proj(q_c)[0]
            else:
                q = self.q_proj(hidden_states)[0]
            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)
688
689
690
691
            if envs.VLLM_USE_LIGHTOP:
                kv_c_normed = self.kv_a_layernorm.forward_cuda_opt(kv_c)
            else:
                kv_c_normed = self.kv_a_layernorm(kv_c.contiguous())
692

693
694
695
            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)
696

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

700
701
702
703
704
705
706
            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]
707
708


wangding zeng's avatar
wangding zeng committed
709
710
711
712
713
class DeepseekV2DecoderLayer(nn.Module):

    def __init__(
        self,
        config: PretrainedConfig,
714
        prefix: str,
715
        model_config: ModelConfig,
wangding zeng's avatar
wangding zeng committed
716
717
        cache_config: Optional[CacheConfig] = None,
        quant_config: Optional[QuantizationConfig] = None,
718
        enable_eplb: bool = False,
wangding zeng's avatar
wangding zeng committed
719
720
721
722
723
724
725
    ) -> 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)
726
727
728
        # DecoderLayers are created with `make_layers` which passes the prefix
        # with the layer's index.
        layer_idx = int(prefix.split(sep='.')[-1])
729
        self.layer_idx = layer_idx
730
731
732
733
734
        if model_config.use_mla:
            attn_cls = DeepseekV2MLAAttention
        else:
            attn_cls = DeepseekV2Attention
        self.self_attn = attn_cls(
wangding zeng's avatar
wangding zeng committed
735
736
737
738
739
740
741
742
743
744
745
746
747
748
            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,
749
            prefix=f"{prefix}.self_attn",
wangding zeng's avatar
wangding zeng committed
750
        )
751

wangding zeng's avatar
wangding zeng committed
752
753
754
        if (config.n_routed_experts is not None
                and layer_idx >= config.first_k_dense_replace
                and layer_idx % config.moe_layer_freq == 0):
755
756
757
758
            self.mlp = DeepseekV2MoE(
                config=config,
                quant_config=quant_config,
                prefix=f"{prefix}.mlp",
759
                enable_eplb=enable_eplb,
760
            )
wangding zeng's avatar
wangding zeng committed
761
762
763
764
765
766
        else:
            self.mlp = DeepseekV2MLP(
                hidden_size=config.hidden_size,
                intermediate_size=config.intermediate_size,
                hidden_act=config.hidden_act,
                quant_config=quant_config,
767
                prefix=f"{prefix}.mlp",
wangding zeng's avatar
wangding zeng committed
768
769
770
771
772
            )
        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)
773
        self.routed_scaling_factor = config.routed_scaling_factor
774
775
        self.use_fused_rms_quant = envs.USE_FUSED_RMS_QUANT
        self.use_fused_custom_all_reduce = envs.USE_FUSED_CUSTOM_ALL_REDUCE_RMS_QUANT
wangding zeng's avatar
wangding zeng committed
776

777
    def forward_fused_rmsquant(
wangding zeng's avatar
wangding zeng committed
778
779
780
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
        residual: Optional[torch.Tensor]
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        # Fix residual FP16 overflow
        residual_fix_overflow = False
        
        assert self.input_layernorm.has_weight is True
        if residual is None:
            residual = hidden_states
            hidden_states, _ = self.self_attn(
                positions = positions,
                hidden_states = hidden_states,
                rms_weight = self.input_layernorm.weight.data,
                residual = None
            )
            residual_fix_overflow = True
        else:
            hidden_states, new_residual = self.self_attn(
                positions = positions,
                hidden_states = hidden_states,
                rms_weight = self.input_layernorm.weight.data,
                residual = residual
            )
            residual = new_residual
804
            
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
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
898
899
900
901
902
903
904
905
906
907
908
909
        if hidden_states.dtype == torch.float16:
            # rmsnorm, and rmsnorm result would not affect by scale.
            hidden_states *= 1. / self.routed_scaling_factor
            if self.layer_idx == 0 or residual_fix_overflow:
                # The residual is shared by all layers, we only scale it on
                # first layer.
                residual *= 1. / self.routed_scaling_factor

        hidden_states, new_resi = self.mlp(hidden_states, self.post_attention_layernorm.weight.data, residual)

        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
            hidden_states *= 1. / self.routed_scaling_factor
        return hidden_states, new_resi

    def forward_fused_CRQ(
        self, 
        positions: torch.Tensor, 
        hidden_states: torch.Tensor, 
        residual: Optional[torch.Tensor]
        ) -> Tuple[torch.Tensor, torch.Tensor]:
        residual_fix_overflow = False
        if residual is None:
            residual = hidden_states
            hidden_states = self.input_layernorm(hidden_states)
            residual_fix_overflow = True
        else:
            hidden_states, resi_new = self.input_layernorm(
                hidden_states, residual)
            residual = resi_new 
        new_hs, new_resi, xq, xs = self.self_attn(
            positions=positions,
            hidden_states=hidden_states,
            pa_rms_weight=self.post_attention_layernorm.weight.data, 
            pa_residual=residual,
            pa_rms_eps=self.post_attention_layernorm.variance_epsilon,
            pa_quant_dtype = torch.int8,
            update_input=True
        )
        
        
        assert xq is not None and xs is not None
        if new_hs.dtype == torch.float16: # overflow处理逻辑
            new_hs *= 1. / self.routed_scaling_factor
            if self.layer_idx == 0 or residual_fix_overflow:
                new_resi *= 1. / self.routed_scaling_factor
            
        hidden_states = self.mlp(new_hs, xqxs=(xq, xs))

        if isinstance(self.mlp,
                    DeepseekV2MLP) and hidden_states.dtype == torch.float16:
            hidden_states *= 1. / self.routed_scaling_factor
        return hidden_states, new_resi
    
    def forward_default(
        self, 
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        residual: Optional[torch.Tensor]
        ) -> Tuple[torch.Tensor, torch.Tensor]:
        # Self Attention
        # Fix residual FP16 overflow
        residual_fix_overflow = False
        if residual is None:
            residual = hidden_states
            hidden_states = self.input_layernorm(hidden_states)
            residual_fix_overflow = True
        else:
            hidden_states, residual = self.input_layernorm(
                hidden_states, residual)
        hidden_states = self.self_attn(
            positions=positions,
            hidden_states=hidden_states,
        )

        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 or residual_fix_overflow:
                # The residual is shared by all layers, we only scale it on
                # first layer.
                residual *= 1. / self.routed_scaling_factor

        # Fully Connected
        hidden_states, residual = self.post_attention_layernorm(
            hidden_states, residual)
        hidden_states = self.mlp(hidden_states)

        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
            hidden_states *= 1. / self.routed_scaling_factor

        return hidden_states, residual
910
                
911
912
913
    def choose_forward(self):
        if self.use_fused_rms_quant:
            return self.forward_fused_rmsquant
914

915
916
        elif self.use_fused_custom_all_reduce:
            return self.forward_fused_CRQ
wangding zeng's avatar
wangding zeng committed
917
        else:
918
            return self.forward_default
919

920
921
922
923
924
925
926
927
928
929
    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        residual: Optional[torch.Tensor]
    )  -> Tuple[torch.Tensor, torch.Tensor]:
        forward_func = self.choose_forward()
        return forward_func(positions=positions, hidden_states=hidden_states, residual=residual )

           
wangding zeng's avatar
wangding zeng committed
930
931


932
@support_torch_compile
wangding zeng's avatar
wangding zeng committed
933
934
935
936
class DeepseekV2Model(nn.Module):

    fall_back_to_pt_during_load = False

937
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
wangding zeng's avatar
wangding zeng committed
938
        super().__init__()
939
940

        config = vllm_config.model_config.hf_config
941
        model_config = vllm_config.model_config
942
943
        cache_config = vllm_config.cache_config
        quant_config = vllm_config.quant_config
944
        enable_eplb = vllm_config.parallel_config.enable_eplb
945
        self.config = config
946

wangding zeng's avatar
wangding zeng committed
947
948
        self.vocab_size = config.vocab_size

949
950
951
952
        if get_pp_group().is_first_rank:
            self.embed_tokens = VocabParallelEmbedding(
                config.vocab_size,
                config.hidden_size,
953
954
                quant_config=quant_config,
                prefix=f"{prefix}.embed_tokens")
955
956
957
958
959
960
961
962
        else:
            self.embed_tokens = PPMissingLayer()

        self.start_layer, self.end_layer, self.layers = make_layers(
            config.num_hidden_layers,
            lambda prefix: DeepseekV2DecoderLayer(
                config,
                prefix,
963
                model_config=model_config,
964
965
                cache_config=cache_config,
                quant_config=quant_config,
966
                enable_eplb=enable_eplb,
967
968
969
970
971
972
973
            ),
            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()
974
975
976
        self.make_empty_intermediate_tensors = (
            make_empty_intermediate_tensors_factory(
                ["hidden_states", "residual"], config.hidden_size))
wangding zeng's avatar
wangding zeng committed
977

978
979
980
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embed_tokens(input_ids)

wangding zeng's avatar
wangding zeng committed
981
982
983
984
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
985
        intermediate_tensors: Optional[IntermediateTensors],
986
        inputs_embeds: Optional[torch.Tensor] = None,
987
    ) -> Union[torch.Tensor, IntermediateTensors]:
988
        if get_pp_group().is_first_rank:
989
990
991
992
            if inputs_embeds is not None:
                hidden_states = inputs_embeds
            else:
                hidden_states = self.get_input_embeddings(input_ids)
993
994
995
996
997
998
            residual = None
        else:
            assert intermediate_tensors is not None
            hidden_states = intermediate_tensors["hidden_states"]
            residual = intermediate_tensors["residual"]

999
        for layer in self.layers[self.start_layer:self.end_layer]:
1000
            hidden_states, residual = layer(positions, hidden_states, residual)
1001
1002
1003
1004
1005
1006
1007

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

wangding zeng's avatar
wangding zeng committed
1008
1009
1010
1011
        hidden_states, _ = self.norm(hidden_states, residual)
        return hidden_states


1012
class DeepseekV2ForCausalLM(nn.Module, SupportsPP, MixtureOfExperts):
wangding zeng's avatar
wangding zeng committed
1013

1014
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
wangding zeng's avatar
wangding zeng committed
1015
        super().__init__()
1016
1017
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
1018
1019
1020
1021
1022
1023
1024

        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'

1025
        self.use_w4a16_moe_sz = os.environ.get('AWQ_MOE_SZ') == '1'
wangding zeng's avatar
wangding zeng committed
1026
1027
        self.config = config
        self.quant_config = quant_config
王敏's avatar
王敏 committed
1028

1029
        self.model = DeepseekV2Model(vllm_config=vllm_config,
1030
                                     prefix=maybe_prefix(prefix, "model"))
1031
1032
1033
1034
1035
1036
        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
1037
        self.logits_processor = LogitsProcessor(config.vocab_size)
1038
1039
        self.make_empty_intermediate_tensors = (
            self.model.make_empty_intermediate_tensors)
1040
1041
1042
1043
1044
1045
1046
1047
        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] = []
1048
        example_moe = None
1049
        for layer in self.model.layers:
1050
1051
1052
            if isinstance(layer, PPMissingLayer):
                continue

1053
1054
            assert isinstance(layer, DeepseekV2DecoderLayer)
            if isinstance(layer.mlp, DeepseekV2MoE):
1055
                example_moe = layer.mlp
1056
1057
1058
1059
1060
1061
1062
1063
1064
                self.moe_layers.append(layer.mlp.experts)

        # Pick last one layer since the first ones may be dense layers.
        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
zhuwenwen's avatar
zhuwenwen committed
1065
        
王敏's avatar
王敏 committed
1066
        self.use_llama_nn = os.environ.get('LLAMA_NN') == '1'
1067
        self.use_awq_pad = os.environ.get('AWQ_PAD') == '1'
1068
1069
        self.tritonsingleton= W8a8GetCacheJSON() 
        self.tritonsingleton.topk = config.num_experts_per_tok
王敏's avatar
王敏 committed
1070
        self.tritonsingleton.quant_method=self.quant_method
王敏's avatar
王敏 committed
1071

1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
    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,
            )
1087

1088
1089
1090
    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
1091
1092
1093
1094
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
1095
        intermediate_tensors: Optional[IntermediateTensors] = None,
1096
        inputs_embeds: Optional[torch.Tensor] = None,
1097
    ) -> Union[torch.Tensor, IntermediateTensors]:
1098
        hidden_states = self.model(input_ids, positions, intermediate_tensors,
1099
                                   inputs_embeds)
wangding zeng's avatar
wangding zeng committed
1100
1101
        return hidden_states

1102
1103
1104
1105
1106
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
1107
        logits = self.logits_processor(self.lm_head, hidden_states,
wangding zeng's avatar
wangding zeng committed
1108
1109
1110
                                       sampling_metadata)
        return logits

1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
    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),
        })
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
        
    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
1144

1145
1146
    def load_weights(self, weights: Iterable[tuple[str,
                                                   torch.Tensor]]) -> set[str]:
wangding zeng's avatar
wangding zeng committed
1147
1148
1149
1150
1151
1152
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
            ("gate_up_proj", "gate_proj", 0),
            ("gate_up_proj", "up_proj", 1),
        ]

1153
1154
        # Params for weights, fp8 weight scales, fp8 activation scales
        # (param_name, weight_name, expert_id, shard_id)
王敏's avatar
王敏 committed
1155
1156
1157
1158
        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",
1159
1160
            num_experts=self.config.n_routed_experts,
            num_redundant_experts=self.num_redundant_experts)
1161

wangding zeng's avatar
wangding zeng committed
1162
        params_dict = dict(self.named_parameters())
1163
        loaded_params: set[str] = set()
wangding zeng's avatar
wangding zeng committed
1164
1165
1166
        for name, loaded_weight in weights:
            if "rotary_emb.inv_freq" in name:
                continue
1167

1168
1169
1170
            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
1171

wangding zeng's avatar
wangding zeng committed
1172
            for (param_name, weight_name, shard_id) in stacked_params_mapping:
1173
                # Skip non-stacked layers and experts (experts handled below).
wangding zeng's avatar
wangding zeng committed
1174
1175
                if weight_name not in name:
                    continue
1176
1177
1178
1179
1180
1181
1182
1183
                # 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
1184
                name = name.replace(weight_name, param_name)
王敏's avatar
王敏 committed
1185

wangding zeng's avatar
wangding zeng committed
1186
1187
1188
                # Skip loading extra bias for GPTQ models.
                if name.endswith(".bias") and name not in params_dict:
                    continue
1189
1190
1191
1192

                if is_pp_missing_parameter(name, self):
                    continue

wangding zeng's avatar
wangding zeng committed
1193
1194
1195
1196
1197
                param = params_dict[name]
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                break
            else:
1198
                is_expert_weight = False
1199
1200
1201
1202
                for mapping in expert_params_mapping:
                    param_name, weight_name, expert_id, shard_id = mapping
                    if weight_name not in name:
                        continue
1203

1204
1205
1206
1207
1208
1209
1210
1211
1212
                    # 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):
1213
1214
                        continue

1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
                    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:
1228
                        name = name_mapped
1229
                        break
1230
                else:
1231
1232
1233
1234
1235
                    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
王敏's avatar
王敏 committed
1236
                    
1237
1238
1239
1240
                    # Skip loading extra bias for GPTQ models.
                    if name.endswith(".bias") and name not in params_dict:
                        continue

1241
1242
1243
1244
1245
                    # Remapping the name of FP8 kv-scale.
                    name = maybe_remap_kv_scale_name(name, params_dict)
                    if name is None:
                        continue

1246
1247
1248
                    if is_pp_missing_parameter(name, self):
                        continue

zhuwenwen's avatar
zhuwenwen committed
1249
1250
1251
1252
                    try:
                        param = params_dict[name]
                    except Exception as e:
                        continue
1253
1254
1255
                    weight_loader = getattr(param, "weight_loader",
                                            default_weight_loader)
                    weight_loader(param, loaded_weight)
1256
            loaded_params.add(name)
王敏's avatar
王敏 committed
1257
            
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
        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
1287
            
1288
        return loaded_params
1289
1290
1291
1292


class DeepseekV3ForCausalLM(DeepseekV2ForCausalLM):
    pass
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304


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