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

wangding zeng's avatar
wangding zeng committed
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
# Adapted from
# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/llama/modeling_llama.py
# Copyright 2023 The vLLM team.
# Copyright 2023 DeepSeek-AI and the HuggingFace Inc. team. All rights reserved.
#
# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
# and OPT implementations in this library. It has been modified from its
# original forms to accommodate minor architectural differences compared
# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
25
"""Inference-only DeepseekV2/DeepseekV3 model."""
王敏'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 itertools import islice
33
from typing import Any, Optional, Union
wangding zeng's avatar
wangding zeng committed
34
35
36

import torch
from torch import nn
37
from transformers import DeepseekV2Config, DeepseekV3Config
wangding zeng's avatar
wangding zeng committed
38

39
import vllm.envs as envs
40
from vllm.attention import Attention
41
from vllm.compilation.decorators import support_torch_compile
42
from vllm.config import CacheConfig, ParallelConfig, VllmConfig
43
from vllm.distributed import (get_ep_group, get_pp_group,
44
45
46
                              get_tensor_model_parallel_rank,
                              get_tensor_model_parallel_world_size,
                              tensor_model_parallel_all_gather)
wangding zeng's avatar
wangding zeng committed
47
from vllm.model_executor.layers.activation import SiluAndMul
48
from vllm.model_executor.layers.fused_moe import FusedMoE
wangding zeng's avatar
wangding zeng committed
49
50
51
52
53
54
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
55
from vllm.model_executor.layers.mla import MLAModules, MultiHeadLatentAttention
56
from vllm.model_executor.layers.quantization import QuantizationConfig
wangding zeng's avatar
wangding zeng committed
57
from vllm.model_executor.layers.rotary_embedding import get_rope
58
from vllm.model_executor.layers.shared_fused_moe import SharedFusedMoE
wangding zeng's avatar
wangding zeng committed
59
60
from vllm.model_executor.layers.vocab_parallel_embedding import (
    ParallelLMHead, VocabParallelEmbedding)
61
62
from vllm.model_executor.model_loader.weight_utils import (
    default_weight_loader, maybe_remap_kv_scale_name)
wangding zeng's avatar
wangding zeng committed
63
from vllm.model_executor.sampling_metadata import SamplingMetadata
64
from vllm.platforms import current_platform
65
from vllm.sequence import IntermediateTensors
66
from vllm.utils import cdiv, direct_register_custom_op
wangding zeng's avatar
wangding zeng committed
67

68
from .interfaces import MixtureOfExperts, SupportsLoRA, SupportsPP
69
from .utils import (PPMissingLayer, is_pp_missing_parameter,
70
71
                    make_empty_intermediate_tensors_factory, make_layers,
                    maybe_prefix)
王敏's avatar
王敏 committed
72
from vllm import _custom_ops as ops
73
from vllm.utils import W8a8GetCacheJSON
74

75
os.environ['DPSK_FP16_QUICK'] = os.environ.get('DPSK_FP16_QUICK', '0')
wangding zeng's avatar
wangding zeng committed
76
77
78
79
80
81
82
83
84
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,
85
        is_sequence_parallel=False,
86
        prefix: str = "",
wangding zeng's avatar
wangding zeng committed
87
88
    ) -> None:
        super().__init__()
89
90
91
92
93

        # If is_sequence_parallel, the input and output tensors are sharded
        # across the ranks within the tp_group. In this case the weights are
        # replicated and no collective ops are needed.
        # Otherwise we use standard TP with an allreduce at the end.
wangding zeng's avatar
wangding zeng committed
94
95
96
        self.gate_up_proj = MergedColumnParallelLinear(
            hidden_size, [intermediate_size] * 2,
            bias=False,
97
            quant_config=quant_config,
98
            disable_tp=is_sequence_parallel,
99
            prefix=f"{prefix}.gate_up_proj")
wangding zeng's avatar
wangding zeng committed
100
101
102
103
        self.down_proj = RowParallelLinear(intermediate_size,
                                           hidden_size,
                                           bias=False,
                                           quant_config=quant_config,
104
                                           reduce_results=reduce_results,
105
                                           disable_tp=is_sequence_parallel,
106
                                           prefix=f"{prefix}.down_proj")
wangding zeng's avatar
wangding zeng committed
107
108
109
110
111
        if hidden_act != "silu":
            raise ValueError(f"Unsupported activation: {hidden_act}. "
                             "Only silu is supported for now.")
        self.act_fn = SiluAndMul()

112
113
114
115
116
117
118
    def forward(self, x,
                rms_weight: Optional[torch.Tensor] = None,
                residual: Optional[torch.Tensor] = None,
                update_hd: Optional[bool] = False
                ):
        if envs.USE_FUSED_RMS_QUANT:
            gate_up, new_resi, _  = self.gate_up_proj(x, rms_weight, residual, update_hd=update_hd)
119
120
121
122
123
            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)
124
125
126
127
128
129
            return x, new_resi
        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
130
131


132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
# Chunk x along the num_tokens axis for sequence parallelism
# NOTE: This is wrapped in a torch custom op to work around the following issue:
# The output tensor can have a sequence length 0 at small input sequence lengths
# even though we explicitly pad to avoid this.
def sequence_parallel_chunk(x: torch.Tensor) -> torch.Tensor:
    tp_size = get_tensor_model_parallel_world_size()
    tp_rank = get_tensor_model_parallel_rank()

    # all_gather needs the sequence length to be divisible by tp_size
    seq_len = x.size(0)
    remainder = seq_len % tp_size
    if remainder != 0:
        pad_len = tp_size - remainder
        x = nn.functional.pad(x, (0, 0, 0, pad_len))

    chunk = x.shape[0] // tp_size
    start = tp_rank * chunk
    return torch.narrow(x, 0, start, chunk)


def sequence_parallel_chunk_fake(x: torch.Tensor) -> torch.Tensor:
    tp_size = get_tensor_model_parallel_world_size()
    seq_len = cdiv(x.size(0), tp_size)
    shape = list(x.shape)
    shape[0] = seq_len
    out = torch.empty(shape, dtype=x.dtype, device=x.device)
    return out


direct_register_custom_op(
    op_name="sequence_parallel_chunk",
    op_func=sequence_parallel_chunk,
    mutates_args=[],
    fake_impl=sequence_parallel_chunk_fake,
    dispatch_key=current_platform.dispatch_key,
    tags=(torch.Tag.needs_fixed_stride_order, ),
)


wangding zeng's avatar
wangding zeng committed
171
172
173
174
class DeepseekV2MoE(nn.Module):

    def __init__(
        self,
175
        config: Union[DeepseekV2Config, DeepseekV3Config],
176
        parallel_config: ParallelConfig,
wangding zeng's avatar
wangding zeng committed
177
        quant_config: Optional[QuantizationConfig] = None,
178
        prefix: str = "",
wangding zeng's avatar
wangding zeng committed
179
180
181
    ):
        super().__init__()
        self.tp_size = get_tensor_model_parallel_world_size()
182
183
        self.tp_rank = get_tensor_model_parallel_rank()

wangding zeng's avatar
wangding zeng committed
184
        self.routed_scaling_factor = config.routed_scaling_factor
185
186
187
188
189
190

        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
191

192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
        # The all_reduce at the end of attention (during o_proj) means that
        # inputs are replicated across each rank of the tensor parallel group.
        # If using expert-parallelism with DeepEP All2All ops, replicated
        # tokens results in useless duplicate computation and communication.
        #
        # In this case, ensure the input to the experts is sequence parallel
        # to avoid the excess work.
        #
        # Not needed for pplx-kernels as it can handle duplicate input tokens.
        self.is_sequence_parallel = (envs.VLLM_ALL2ALL_BACKEND
                                     in ("deepep_high_throughput",
                                         "deepep_low_latency")
                                     and parallel_config.enable_expert_parallel
                                     and self.tp_size > 1)

207
208
209
210
        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
211
        self.gate = ReplicatedLinear(config.hidden_size,
212
                                     config.n_routed_experts,
wangding zeng's avatar
wangding zeng committed
213
                                     bias=False,
214
215
                                     quant_config=None,
                                     prefix=f"{prefix}.gate")
216
217
        if config.topk_method == "noaux_tc":
            self.gate.e_score_correction_bias = nn.Parameter(
218
                torch.empty(config.n_routed_experts, dtype=torch.float32))
219
220
221
        else:
            self.gate.e_score_correction_bias = None

222
        # Load balancing settings.
223
224
        eplb_config = parallel_config.eplb_config
        self.enable_eplb = parallel_config.enable_eplb
225
        self.dpsk_fp16_quick = os.environ.get('DPSK_FP16_QUICK') == '1'
226

227
        self.n_redundant_experts = eplb_config.num_redundant_experts
228
229
230
231
232
233
234
235
236
237
        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)

238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
        if config.n_shared_experts is 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,
                # we do scaling outside, set factor to 1.0 to avoid double mul
                routed_scaling_factor=1.0,
                e_score_correction_bias=self.gate.e_score_correction_bias,
                enable_eplb=self.enable_eplb,
256
257
258
                num_redundant_experts=self.n_redundant_experts,
                is_sequence_parallel=self.is_sequence_parallel,
            )
259
260
            self.shared_experts = None
        else:
wangding zeng's avatar
wangding zeng committed
261
262
            intermediate_size = (config.moe_intermediate_size *
                                 config.n_shared_experts)
263

wangding zeng's avatar
wangding zeng committed
264
265
266
267
268
            self.shared_experts = DeepseekV2MLP(
                hidden_size=config.hidden_size,
                intermediate_size=intermediate_size,
                hidden_act=config.hidden_act,
                quant_config=quant_config,
269
                is_sequence_parallel=self.is_sequence_parallel,
270
                reduce_results=False,
271
                prefix=f"{prefix}.shared_experts",
wangding zeng's avatar
wangding zeng committed
272
            )
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
            self.experts = SharedFusedMoE(
                shared_experts=self.shared_experts,
                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,
                # we do scaling outside, set factor to 1.0 to avoid double mul
                routed_scaling_factor=1.0,
                e_score_correction_bias=self.gate.e_score_correction_bias,
                enable_eplb=self.enable_eplb,
291
292
293
                num_redundant_experts=self.n_redundant_experts,
                is_sequence_parallel=self.is_sequence_parallel,
            )
zhuwenwen's avatar
zhuwenwen committed
294
295
296
        from vllm.two_batch_overlap.two_batch_overlap import tbo_all_reduce
        self.tbo_all_reduce = tbo_all_reduce

297

298
299
300
301
    def forward(self, hidden_states: torch.Tensor,
                rms_weight: Optional[torch.Tensor] = None,
                residual: Optional[torch.Tensor] = None
                ) -> torch.Tensor:
wangding zeng's avatar
wangding zeng committed
302
303
        num_tokens, hidden_dim = hidden_states.shape
        hidden_states = hidden_states.view(-1, hidden_dim)
304
305
306
307
308
309
310
311
312

        # Chunk the hidden states so they aren't replicated across TP ranks.
        # This avoids duplicate computation in self.experts.
        # TODO: We can replace the all_reduce at the end of attn with a
        # reduce_scatter instead of chunking here.
        if self.is_sequence_parallel:
            hidden_states = torch.ops.vllm.sequence_parallel_chunk(
                hidden_states)

wangding zeng's avatar
wangding zeng committed
313
314
        # router_logits: (num_tokens, n_experts)
        router_logits, _ = self.gate(hidden_states)
315
        
316
317
318
319
        fused_moe_out = self.experts(hidden_states=hidden_states,
                                     router_logits=router_logits)
        if self.shared_experts is not None:
            shared_output, final_hidden_states = fused_moe_out
320
        else:
321
322
            shared_output = None
            final_hidden_states = fused_moe_out
zhuwenwen's avatar
zhuwenwen committed
323
        
324
325
        # Fix FP16 overflow
        # See DeepseekV2DecoderLayer for more details.
326
        if hidden_states.dtype != torch.float16 or self.dpsk_fp16_quick:
327
328
329
330
331
332
333
334
            final_hidden_states *= self.routed_scaling_factor
        elif self.shared_experts is not None:
            assert shared_output is not None
            shared_output *= (1. / self.routed_scaling_factor)

        if self.shared_experts is not None:
            assert shared_output is not None
            final_hidden_states += shared_output
335

336
337
338
339
340
        if self.is_sequence_parallel:
            final_hidden_states = tensor_model_parallel_all_gather(
                final_hidden_states, 0)
            final_hidden_states = final_hidden_states[:num_tokens]
        elif self.tp_size > 1:
341
            if envs.VLLM_ENABLE_TBO:
lizhigong's avatar
lizhigong committed
342
343
                final_hidden_states = self.tbo_all_reduce(final_hidden_states)
            else:
zhuwenwen's avatar
zhuwenwen committed
344
345
346
                final_hidden_states = (
                    self.experts.maybe_all_reduce_tensor_model_parallel(
                        final_hidden_states))
wangding zeng's avatar
wangding zeng committed
347
348
349
350
351
352
353
354
355
356
357
358
359
360
        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,
361
        config: Union[DeepseekV2Config, DeepseekV3Config],
wangding zeng's avatar
wangding zeng committed
362
363
364
365
366
367
368
369
        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,
370
        rope_scaling: Optional[dict[str, Any]] = None,
wangding zeng's avatar
wangding zeng committed
371
372
373
        max_position_embeddings: int = 8192,
        cache_config: Optional[CacheConfig] = None,
        quant_config: Optional[QuantizationConfig] = None,
374
        prefix: str = "",
wangding zeng's avatar
wangding zeng committed
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
    ) -> 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,
396
397
                                             quant_config=quant_config,
                                             prefix=f"{prefix}.q_a_proj")
wangding zeng's avatar
wangding zeng committed
398
399
400
401
402
403
            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,
404
405
                                                 quant_config=quant_config,
                                                 prefix=f"{prefix}.q_b_proj")
wangding zeng's avatar
wangding zeng committed
406
407
408
409
410
        else:
            self.q_proj = ColumnParallelLinear(self.hidden_size,
                                               self.num_heads *
                                               self.qk_head_dim,
                                               bias=False,
411
412
                                               quant_config=quant_config,
                                               prefix=f"{prefix}.q_proj")
wangding zeng's avatar
wangding zeng committed
413

414
415
416
417
418
419
        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
420
421
422
423
424
425
        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,
426
427
            quant_config=quant_config,
            prefix=f"{prefix}.kv_b_proj")
wangding zeng's avatar
wangding zeng committed
428
429
430
431
        # O projection.
        self.o_proj = RowParallelLinear(self.num_heads * self.v_head_dim,
                                        self.hidden_size,
                                        bias=False,
432
433
                                        quant_config=quant_config,
                                        prefix=f"{prefix}.o_proj")
434
435
        if rope_scaling:
            rope_scaling["rope_type"] = 'deepseek_yarn'
436

wangding zeng's avatar
wangding zeng committed
437
438
439
440
441
442
443
444
445
446
447
448
449
450
        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,
451
                              self.qk_head_dim,
wangding zeng's avatar
wangding zeng committed
452
453
454
                              self.scaling,
                              num_kv_heads=self.num_local_heads,
                              cache_config=cache_config,
455
456
                              quant_config=quant_config,
                              prefix=f"{prefix}.attn")
wangding zeng's avatar
wangding zeng committed
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476

    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)
477
        kv_a = self.kv_a_layernorm(kv_a)
wangding zeng's avatar
wangding zeng committed
478
479
480
481
482
        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:]
483

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

wangding zeng's avatar
wangding zeng committed
486
487
488
489
        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
490
491
492
493
        # 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)
494
        attn_output = self.attn(q, k, v)
wangding zeng's avatar
wangding zeng committed
495
        attn_output = attn_output.view(
496
497
            -1, self.num_local_heads,
            self.qk_head_dim)[..., :self.v_head_dim].reshape(
wangding zeng's avatar
wangding zeng committed
498
499
500
501
502
                -1, self.num_local_heads * self.v_head_dim)
        output, _ = self.o_proj(attn_output)
        return output


503
504
505
506
507
508
509
510
511
512
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,
513
        config: Union[DeepseekV2Config, DeepseekV3Config],
514
515
516
517
518
519
520
521
        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,
522
        rope_scaling: Optional[dict[str, Any]] = None,
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
        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:
548
            self.fused_qkv_a_proj = MergedColumnParallelLinear(
549
550
551
552
                self.hidden_size,
                [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
                bias=False,
                quant_config=quant_config,
553
554
                prefix=f"{prefix}.fused_qkv_a_proj",
                disable_tp=True)
555
556
557
558
559
560
561
562
563
        else:
            self.kv_a_proj_with_mqa = ReplicatedLinear(
                self.hidden_size,
                self.kv_lora_rank + self.qk_rope_head_dim,
                bias=False,
                quant_config=quant_config,
                prefix=f"{prefix}.kv_a_proj_with_mqa")

        if self.q_lora_rank is not None:
564
565
            self.q_a_layernorm = RMSNorm(self.q_lora_rank,
                                         eps=config.rms_norm_eps)
566
            self.q_b_proj = ColumnParallelLinear(self.q_lora_rank,
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
                                                 self.num_heads *
                                                 self.qk_head_dim,
                                                 bias=False,
                                                 quant_config=quant_config,
                                                 prefix=f"{prefix}.q_b_proj")
        else:
            self.q_proj = ColumnParallelLinear(self.hidden_size,
                                               self.num_heads *
                                               self.qk_head_dim,
                                               bias=False,
                                               quant_config=quant_config,
                                               prefix=f"{prefix}.q_proj")
        self.kv_a_layernorm = RMSNorm(self.kv_lora_rank,
                                      eps=config.rms_norm_eps)
        self.kv_b_proj = ColumnParallelLinear(
            self.kv_lora_rank,
            self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.kv_b_proj")
        self.o_proj = RowParallelLinear(self.num_heads * self.v_head_dim,
                                        self.hidden_size,
                                        bias=False,
                                        quant_config=quant_config,
                                        prefix=f"{prefix}.o_proj")

593
594
        if rope_scaling:
            rope_scaling["rope_type"] = 'deepseek_yarn'
595
596
597
598
599
600
601
602
603
604
605
606
        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

607
608
        mla_modules = MLAModules(
            kv_a_layernorm=self.kv_a_layernorm,
609
            kv_b_proj=self.kv_b_proj,
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
            rotary_emb=self.rotary_emb,
            o_proj=self.o_proj,
            fused_qkv_a_proj=self.fused_qkv_a_proj
            if self.q_lora_rank is not None else None,
            kv_a_proj_with_mqa=self.kv_a_proj_with_mqa
            if self.q_lora_rank is None else None,
            q_a_layernorm=self.q_a_layernorm
            if self.q_lora_rank is not None else None,
            q_b_proj=self.q_b_proj if self.q_lora_rank is not None else None,
            q_proj=self.q_proj if self.q_lora_rank is None else None,
        )
        self.mla_attn = MultiHeadLatentAttention(
            self.hidden_size,
            self.num_local_heads,
            self.scaling,
            self.qk_nope_head_dim,
            self.qk_rope_head_dim,
            self.v_head_dim,
            self.q_lora_rank,
            self.kv_lora_rank,
            mla_modules,
            cache_config,
            quant_config,
            prefix,
634
635
636
637
638
639
640
        )

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
641
        return self.mla_attn(positions, hidden_states)
642
643


wangding zeng's avatar
wangding zeng committed
644
645
class DeepseekV2DecoderLayer(nn.Module):

646
    def __init__(self, vllm_config: VllmConfig, prefix: str) -> None:
wangding zeng's avatar
wangding zeng committed
647
        super().__init__()
648
649
650
651
652
653
654

        config = vllm_config.model_config.hf_config
        model_config = vllm_config.model_config
        cache_config = vllm_config.cache_config
        quant_config = vllm_config.quant_config
        parallel_config = vllm_config.parallel_config

wangding zeng's avatar
wangding zeng committed
655
656
657
658
659
        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)
660
661
662
        # DecoderLayers are created with `make_layers` which passes the prefix
        # with the layer's index.
        layer_idx = int(prefix.split(sep='.')[-1])
663
        self.layer_idx = layer_idx
664
665
666
667
668
        if model_config.use_mla:
            attn_cls = DeepseekV2MLAAttention
        else:
            attn_cls = DeepseekV2Attention
        self.self_attn = attn_cls(
wangding zeng's avatar
wangding zeng committed
669
670
671
672
673
674
675
676
677
678
679
680
681
682
            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,
683
            prefix=f"{prefix}.self_attn",
wangding zeng's avatar
wangding zeng committed
684
        )
685
        self.dpsk_fp16_quick = os.environ.get('DPSK_FP16_QUICK') == '1'
wangding zeng's avatar
wangding zeng committed
686
687
688
        if (config.n_routed_experts is not None
                and layer_idx >= config.first_k_dense_replace
                and layer_idx % config.moe_layer_freq == 0):
689
690
            self.mlp = DeepseekV2MoE(
                config=config,
691
                parallel_config=parallel_config,
692
693
694
                quant_config=quant_config,
                prefix=f"{prefix}.mlp",
            )
wangding zeng's avatar
wangding zeng committed
695
696
697
698
699
700
        else:
            self.mlp = DeepseekV2MLP(
                hidden_size=config.hidden_size,
                intermediate_size=config.intermediate_size,
                hidden_act=config.hidden_act,
                quant_config=quant_config,
701
                prefix=f"{prefix}.mlp",
wangding zeng's avatar
wangding zeng committed
702
703
704
705
706
            )
        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)
707
        self.routed_scaling_factor = config.routed_scaling_factor
wangding zeng's avatar
wangding zeng committed
708
709
710
711
712
713
714

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        residual: Optional[torch.Tensor],
    ) -> torch.Tensor:
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
        if envs.USE_FUSED_RMS_QUANT:
            # 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
                
            if hidden_states.dtype == torch.float16 and not self.dpsk_fp16_quick:
                # 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 and not self.dpsk_fp16_quick:
                # 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

wangding zeng's avatar
wangding zeng committed
758
        else:
759
760
761
762
763
764
765
766
767
768
769
770
771
772
            # 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,
            )
wangding zeng's avatar
wangding zeng committed
773

774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
            if hidden_states.dtype == torch.float16 and not self.dpsk_fp16_quick:
                # 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 and not self.dpsk_fp16_quick:
                # 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
wangding zeng's avatar
wangding zeng committed
797
798


799
@support_torch_compile
wangding zeng's avatar
wangding zeng committed
800
801
802
803
class DeepseekV2Model(nn.Module):

    fall_back_to_pt_during_load = False

804
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
wangding zeng's avatar
wangding zeng committed
805
        super().__init__()
806
807
808

        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
809
        self.config = config
810

wangding zeng's avatar
wangding zeng committed
811
812
        self.vocab_size = config.vocab_size

813
814
815
816
        if get_pp_group().is_first_rank:
            self.embed_tokens = VocabParallelEmbedding(
                config.vocab_size,
                config.hidden_size,
817
818
                quant_config=quant_config,
                prefix=f"{prefix}.embed_tokens")
819
820
821
822
823
        else:
            self.embed_tokens = PPMissingLayer()

        self.start_layer, self.end_layer, self.layers = make_layers(
            config.num_hidden_layers,
824
            lambda prefix: DeepseekV2DecoderLayer(vllm_config, prefix),
825
826
827
828
829
830
            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()
831
832
833
        self.make_empty_intermediate_tensors = (
            make_empty_intermediate_tensors_factory(
                ["hidden_states", "residual"], config.hidden_size))
wangding zeng's avatar
wangding zeng committed
834

835
836
837
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embed_tokens(input_ids)

wangding zeng's avatar
wangding zeng committed
838
839
840
841
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
842
        intermediate_tensors: Optional[IntermediateTensors],
843
        inputs_embeds: Optional[torch.Tensor] = None,
844
    ) -> Union[torch.Tensor, IntermediateTensors]:
845
        if get_pp_group().is_first_rank:
846
847
848
849
            if inputs_embeds is not None:
                hidden_states = inputs_embeds
            else:
                hidden_states = self.get_input_embeddings(input_ids)
850
851
852
853
854
855
            residual = None
        else:
            assert intermediate_tensors is not None
            hidden_states = intermediate_tensors["hidden_states"]
            residual = intermediate_tensors["residual"]

856
        for layer in islice(self.layers, self.start_layer, self.end_layer):
857
            hidden_states, residual = layer(positions, hidden_states, residual)
858
859
860
861
862
863
864

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

wangding zeng's avatar
wangding zeng committed
865
866
867
868
        hidden_states, _ = self.norm(hidden_states, residual)
        return hidden_states


869
870
class DeepseekV2ForCausalLM(nn.Module, SupportsPP, MixtureOfExperts,
                            SupportsLoRA):
871
872
873
    packed_modules_mapping = {
        "gate_up_proj": ["gate_proj", "up_proj"],
    }
wangding zeng's avatar
wangding zeng committed
874

875
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
wangding zeng's avatar
wangding zeng committed
876
        super().__init__()
877
878
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
879
880
881
882
883
884
885

        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'

886
        self.use_w4a16_moe_sz = os.environ.get('AWQ_MOE_SZ') == '1'
wangding zeng's avatar
wangding zeng committed
887
888
        self.config = config
        self.quant_config = quant_config
王敏's avatar
王敏 committed
889

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

902

903
        self.model = DeepseekV2Model(vllm_config=vllm_config,
904
                                     prefix=maybe_prefix(prefix, "model"))
905
906
907
908
909
910
        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
911
        self.logits_processor = LogitsProcessor(config.vocab_size)
912
913
        self.make_empty_intermediate_tensors = (
            self.model.make_empty_intermediate_tensors)
914
915
916
917
918
919
920
921
        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] = []
922
        example_moe = None
923
        for layer in self.model.layers:
924
925
926
            if isinstance(layer, PPMissingLayer):
                continue

927
928
            assert isinstance(layer, DeepseekV2DecoderLayer)
            if isinstance(layer.mlp, DeepseekV2MoE):
929
930
                # Pick last one layer since the first ones may be dense layers.
                example_moe = layer.mlp
931
932
                self.moe_layers.append(layer.mlp.experts)

933
934
935
        if example_moe is None:
            raise RuntimeError("No DeepseekV2MoE layer found in model.layers.")

936
937
938
939
940
941
        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
942
        
王敏's avatar
王敏 committed
943
        self.use_llama_nn = os.environ.get('LLAMA_NN') == '1'
944
        self.use_awq_pad = os.environ.get('AWQ_PAD') == '1'
945
946
947
        self.tritonsingleton= W8a8GetCacheJSON() 
        self.tritonsingleton.topk = config.num_experts_per_tok
        self.tritonsingleton.quant_method=self.quant_method 
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963

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

965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
    def update_physical_experts_metadata(
        self,
        num_physical_experts: int,
        num_local_physical_experts: int,
    ) -> None:
        assert self.num_local_physical_experts == num_local_physical_experts
        self.num_physical_experts = num_physical_experts
        self.num_local_physical_experts = num_local_physical_experts
        self.num_redundant_experts = (num_physical_experts -
                                      self.num_logical_experts)
        for layer in self.model.layers:
            if isinstance(layer.mlp, DeepseekV2MoE):
                moe = layer.mlp
                moe.n_local_physical_experts = num_local_physical_experts
                moe.n_physical_experts = num_physical_experts
                moe.n_redundant_experts = self.num_redundant_experts
                moe.experts.update_expert_map()

983
984
985
    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
986
987
988
989
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
990
        intermediate_tensors: Optional[IntermediateTensors] = None,
991
        inputs_embeds: Optional[torch.Tensor] = None,
992
    ) -> Union[torch.Tensor, IntermediateTensors]:
993
        hidden_states = self.model(input_ids, positions, intermediate_tensors,
994
                                   inputs_embeds)
wangding zeng's avatar
wangding zeng committed
995
996
        return hidden_states

997
998
999
1000
1001
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
1002
        logits = self.logits_processor(self.lm_head, hidden_states,
wangding zeng's avatar
wangding zeng committed
1003
1004
1005
                                       sampling_metadata)
        return logits

1006

1007
1008
    def load_weights(self, weights: Iterable[tuple[str,
                                                   torch.Tensor]]) -> set[str]:
wangding zeng's avatar
wangding zeng committed
1009
1010
1011
1012
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
            ("gate_up_proj", "gate_proj", 0),
            ("gate_up_proj", "up_proj", 1),
1013
1014
            ("fused_qkv_a_proj", "q_a_proj", 0),
            ("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1),
wangding zeng's avatar
wangding zeng committed
1015
1016
        ]

1017
1018
        # Params for weights, fp8 weight scales, fp8 activation scales
        # (param_name, weight_name, expert_id, shard_id)
王敏's avatar
王敏 committed
1019
1020
1021
1022
        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",
1023
1024
            num_experts=self.config.n_routed_experts,
            num_redundant_experts=self.num_redundant_experts)
1025

wangding zeng's avatar
wangding zeng committed
1026
        params_dict = dict(self.named_parameters())
1027
        loaded_params: set[str] = set()
wangding zeng's avatar
wangding zeng committed
1028
1029
1030
        for name, loaded_weight in weights:
            if "rotary_emb.inv_freq" in name:
                continue
1031

1032
1033
1034
            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
1035

wangding zeng's avatar
wangding zeng committed
1036
            for (param_name, weight_name, shard_id) in stacked_params_mapping:
1037
                # Skip non-stacked layers and experts (experts handled below).
wangding zeng's avatar
wangding zeng committed
1038
1039
                if weight_name not in name:
                    continue
1040
1041
1042
1043
1044
1045
1046
1047
                # 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
1048
                name_mapped = name.replace(weight_name, param_name)
1049
1050
1051

                # QKV fusion is optional, fall back to normal
                # weight loading if it's not enabled
1052
                # if go with fusion option, then update name
1053
                if ((param_name == "fused_qkv_a_proj")
1054
                        and name_mapped not in params_dict):
1055
                    continue
1056
1057
                else:
                    name = name_mapped
wangding zeng's avatar
wangding zeng committed
1058
1059
1060
                # Skip loading extra bias for GPTQ models.
                if name.endswith(".bias") and name not in params_dict:
                    continue
1061
1062
1063
1064

                if is_pp_missing_parameter(name, self):
                    continue

wangding zeng's avatar
wangding zeng committed
1065
1066
1067
1068
1069
                param = params_dict[name]
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                break
            else:
1070
                is_expert_weight = False
1071
1072
1073
1074
                for mapping in expert_params_mapping:
                    param_name, weight_name, expert_id, shard_id = mapping
                    if weight_name not in name:
                        continue
1075

1076
1077
1078
1079
1080
1081
1082
1083
1084
                    # 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):
1085
1086
                        continue

1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
                    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:
1100
                        name = name_mapped
1101
                        break
1102
                else:
1103
1104
1105
1106
1107
1108
                    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

1109
1110
1111
1112
                    # Skip loading extra bias for GPTQ models.
                    if name.endswith(".bias") and name not in params_dict:
                        continue

1113
1114
1115
1116
1117
                    # Remapping the name of FP8 kv-scale.
                    name = maybe_remap_kv_scale_name(name, params_dict)
                    if name is None:
                        continue

1118
1119
1120
                    if is_pp_missing_parameter(name, self):
                        continue

1121
1122
1123
1124
                    param = params_dict[name]
                    weight_loader = getattr(param, "weight_loader",
                                            default_weight_loader)
                    weight_loader(param, loaded_weight)
1125
            loaded_params.add(name)
王敏's avatar
王敏 committed
1126
            
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
        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
1156
            
1157
        return loaded_params
1158
1159
1160
1161


class DeepseekV3ForCausalLM(DeepseekV2ForCausalLM):
    pass
1162
1163


1164
1165
1166
1167
# Compatibility with
# https://huggingface.co/deepseek-ai/DeepSeek-V3-Base/blob/main/configuration_deepseek.py
def get_spec_layer_idx_from_weight_name(config: Union[DeepseekV2Config,
                                                      DeepseekV3Config],
1168
                                        weight_name: str) -> Optional[int]:
1169
1170
    if (hasattr(config, "num_nextn_predict_layers")
            and config.num_nextn_predict_layers > 0):
1171
1172
1173
1174
1175
        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