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

zhuwenwen's avatar
zhuwenwen committed
4
5
import os
import re
zhuwenwen's avatar
zhuwenwen committed
6

7
8
import typing
from collections.abc import Callable, Iterable
9
10
11
12
13

import torch
import torch.nn as nn
from transformers import PretrainedConfig

14
from vllm._aiter_ops import rocm_aiter_ops
15
from vllm.compilation.decorators import support_torch_compile
16
from vllm.config import VllmConfig
17
from vllm.logger import init_logger
18
from vllm.model_executor.layers.fused_moe import SharedFusedMoE
19
20
21
22
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.logits_processor import LogitsProcessor
from vllm.model_executor.layers.quantization import QuantizationConfig
from vllm.model_executor.layers.vocab_parallel_embedding import (
23
24
25
    ParallelLMHead,
    VocabParallelEmbedding,
)
26
27
28
29
from vllm.model_executor.model_loader.weight_utils import (
    default_weight_loader,
    maybe_remap_kv_scale_name,
)
30
from vllm.platforms import current_platform
31
32
from vllm.sequence import IntermediateTensors

33
34
from .deepseek_v2 import (
    DeepseekV2DecoderLayer,
35
36
    DeepseekV2MixtureOfExperts,
    DeepseekV2MoE,
37
38
    get_spec_layer_idx_from_weight_name,
)
39
from .utils import maybe_prefix
王敏's avatar
王敏 committed
40
from .interfaces import SupportsPP
zhuwenwen's avatar
zhuwenwen committed
41
from vllm import _custom_ops as ops
zhuwenwen's avatar
zhuwenwen committed
42
from vllm.model_executor.layers.quantization.blockwise_int8 import BlockInt8Config
43
import vllm.envs as envs
44
from vllm.utils import direct_register_custom_op
45
logger = init_logger(__name__)
46
47
48
49
50
51


class SharedHead(nn.Module):
    def __init__(
        self,
        config: PretrainedConfig,
52
        prefix: str,
53
        quant_config: QuantizationConfig | None = None,
54
55
56
    ) -> None:
        super().__init__()
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
57
58
59
60
61
62
        self.head = ParallelLMHead(
            config.vocab_size,
            config.hidden_size,
            quant_config=quant_config,
            prefix=maybe_prefix(prefix, "head"),
        )
63
64
65
66
67
68

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self.norm(hidden_states)


class DeepSeekMultiTokenPredictorLayer(nn.Module):
69
    def __init__(self, vllm_config: VllmConfig, prefix: str) -> None:
70
        super().__init__()
71

72
73
        config = vllm_config.speculative_config.draft_model_config.hf_config
        self.config = config
74
75
        quant_config = vllm_config.quant_config

76
77
        self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
78
        self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False)
79

80
        self.device = current_platform.device_type
81

yangql's avatar
yangql committed
82
83
84
        #添加判断,默认开启DSA
        force_disable_dsa = os.environ.get("VLLM_DISABLE_DSA", "0") == "1"
        self.is_v32 = hasattr(config, "index_topk") and not force_disable_dsa
85
86
87
88
89
90
        if self.is_v32:
            topk_tokens = config.index_topk
            topk_indices_buffer = torch.empty(
                vllm_config.scheduler_config.max_num_batched_tokens,
                topk_tokens,
                dtype=torch.int32,
91
                device=self.device,
92
            )
93
94
        else:
            topk_indices_buffer = None
95

96
97
98
99
        self.shared_head = SharedHead(
            config=config, prefix=prefix, quant_config=quant_config
        )
        self.mtp_block = DeepseekV2DecoderLayer(
100
101
102
103
            vllm_config,
            prefix,
            config=self.config,
            topk_indices_buffer=topk_indices_buffer,
104
        )
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
    def fuse_fill_rms_x2_concat(hidden_states_fuse: torch.Tensor, positions: torch.Tensor, inputs_embeds: torch.Tensor,
                                    previous_hidden_states: torch.Tensor, weight_inputs_embeds: torch.Tensor, 
                                    weight_previous_hidden_states: torch.Tensor, epsilon: float) -> None:
        from lightop import fuse_fill_rms_x2_concat
        fuse_fill_rms_x2_concat(hidden_states_fuse, positions, inputs_embeds, previous_hidden_states, weight_inputs_embeds, weight_previous_hidden_states, epsilon)

    def fuse_fill_rms_x2_concat_fake(hidden_states_fuse: torch.Tensor, positions: torch.Tensor, inputs_embeds: torch.Tensor,
                                    previous_hidden_states: torch.Tensor, weight_inputs_embeds: torch.Tensor, 
                                    weight_previous_hidden_states: torch.Tensor, epsilon: float) -> None:
        pass

    direct_register_custom_op(
        op_name="fuse_fill_rms_x2_concat",
        op_func=fuse_fill_rms_x2_concat,
        mutates_args=["hidden_states_fuse", "inputs_embeds"], 
        fake_impl=fuse_fill_rms_x2_concat_fake,
    )

123
124
125
126
127
128

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        previous_hidden_states: torch.Tensor,
129
        inputs_embeds: torch.Tensor | None = None,
130
131
132
133
        spec_step_index: int = 0,
    ) -> torch.Tensor:
        assert inputs_embeds is not None
        # masking inputs at position 0, as not needed by MTP
134
135
136
137
138
139
140
141
        if envs.VLLM_USE_FUSED_FILL_RMS_CAT:
            hidden_states_fuse = torch.empty(inputs_embeds.shape[0], inputs_embeds.shape[1]*2, device=inputs_embeds.device, dtype=inputs_embeds.dtype)
            torch.ops.vllm.fuse_fill_rms_x2_concat(hidden_states_fuse, positions, inputs_embeds, previous_hidden_states, self.enorm.weight, self.hnorm.weight, self.enorm.variance_epsilon)
            hidden_states = self.eh_proj(hidden_states_fuse)
        else:
            inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds)
            inputs_embeds = self.enorm(inputs_embeds)
            previous_hidden_states = self.hnorm(previous_hidden_states)
142

143
144
145
            hidden_states = self.eh_proj(
                torch.cat([inputs_embeds, previous_hidden_states], dim=-1)
            )
146

147
148
149
        hidden_states, residual = self.mtp_block(
            positions=positions, hidden_states=hidden_states, residual=None
        )
150
        hidden_states = residual + hidden_states
151
        return hidden_states
152
153
154
155
156
157
158
159
160


class DeepSeekMultiTokenPredictor(nn.Module):
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        super().__init__()
        config = vllm_config.model_config.hf_config
        self.mtp_start_layer_idx = config.num_hidden_layers
        self.num_mtp_layers = config.num_nextn_predict_layers
        # to map the exact layer index from weights
161

162
163
164
165
166
167
168
169
170
171
172
        self.layers = torch.nn.ModuleDict(
            {
                str(idx): DeepSeekMultiTokenPredictorLayer(
                    vllm_config, f"{prefix}.layers.{idx}"
                )
                for idx in range(
                    self.mtp_start_layer_idx,
                    self.mtp_start_layer_idx + self.num_mtp_layers,
                )
            }
        )
173
174
175
        self.embed_tokens = VocabParallelEmbedding(
            config.vocab_size,
            config.hidden_size,
176
            prefix=maybe_prefix(prefix, "embed_tokens"),
177
        )
178
179
        self.logits_processor = LogitsProcessor(config.vocab_size)

180
    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
181
182
        return self.embed_tokens(input_ids)

183
184
185
186
187
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        previous_hidden_states: torch.Tensor,
188
        inputs_embeds: torch.Tensor | None = None,
189
190
        spec_step_idx: int = 0,
    ) -> torch.Tensor:
191
192
        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)
193
        current_step_idx = spec_step_idx % self.num_mtp_layers
194
        return self.layers[str(self.mtp_start_layer_idx + current_step_idx)](
195
196
197
198
            input_ids,
            positions,
            previous_hidden_states,
            inputs_embeds,
199
            current_step_idx,
200
201
202
203
204
205
206
        )

    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        spec_step_idx: int = 0,
    ) -> torch.Tensor:
207
208
209
210
211
        current_step_idx = spec_step_idx % self.num_mtp_layers
        mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)]
        logits = self.logits_processor(
            mtp_layer.shared_head.head, mtp_layer.shared_head(hidden_states)
        )
212
213
214
        return logits


215
@support_torch_compile
王敏's avatar
王敏 committed
216
class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts, SupportsPP):
217
218
219
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        super().__init__()
        self.config = vllm_config.model_config.hf_config
220
221
222
        self.model = DeepSeekMultiTokenPredictor(
            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
        )
223
224
        # Set MoE hyperparameters
        self.set_moe_parameters()
zhuwenwen's avatar
zhuwenwen committed
225
226
227
228
229
230
231
        quant_config = vllm_config.quant_config

        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'
zhuwenwen's avatar
zhuwenwen committed
232
            # The AWQ layer of MTP uses BlockInt8W8A8.
yangql's avatar
yangql committed
233
            if self.quant_method == "moe_wna16" or self.quant_method == "awq_marlin":
zhuwenwen's avatar
zhuwenwen committed
234
                vllm_config.quant_config = BlockInt8Config(is_checkpoint_int8_serialized=True, weight_block_size=[128,128])
zhuwenwen's avatar
zhuwenwen committed
235
        self.use_llama_nn = os.environ.get('LLAMA_NN') == '1'
236
                
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254

    def set_moe_parameters(self):
        self.expert_weights = []
        self.num_moe_layers = self.config.num_nextn_predict_layers
        self.num_expert_groups = self.config.n_group

        self.moe_layers = []
        self.moe_mlp_layers = []
        example_moe = None
        for layer in self.model.layers.values():
            assert isinstance(layer, DeepSeekMultiTokenPredictorLayer)
            layer = layer.mtp_block
            assert isinstance(layer, DeepseekV2DecoderLayer)
            if isinstance(layer.mlp, DeepseekV2MoE):
                example_moe = layer.mlp
                self.moe_mlp_layers.append(layer.mlp)
                self.moe_layers.append(layer.mlp.experts)
        self.extract_moe_parameters(example_moe)
255

256
257
    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.model.embed_input_ids(input_ids)
258
259
260

    def forward(
        self,
261
        input_ids: torch.Tensor | None,
262
        positions: torch.Tensor,
263
        hidden_states: torch.Tensor,
264
265
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
266
267
        spec_step_idx: int = 0,
    ) -> torch.Tensor:
268
269
270
        hidden_states = self.model(
            input_ids, positions, hidden_states, inputs_embeds, spec_step_idx
        )
271
272
273
274
275
276
        return hidden_states

    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        spec_step_idx: int = 0,
277
    ) -> torch.Tensor | None:
278
        return self.model.compute_logits(hidden_states, spec_step_idx)
279

280
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
281
282
283
        rocm_aiter_moe_shared_expert_enabled = (
            rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
        )
284
285
286
        stacked_params_mapping = [
            ("gate_up_proj", "gate_proj", 0),
            ("gate_up_proj", "up_proj", 1),
287
288
            ("fused_qkv_a_proj", "q_a_proj", 0),
            ("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1),
289
290
        ]

291
        expert_params_mapping = SharedFusedMoE.make_expert_params_mapping(
292
            self,
293
294
295
            ckpt_gate_proj_name="gate_proj",
            ckpt_down_proj_name="down_proj",
            ckpt_up_proj_name="up_proj",
296
297
298
299
300
301
            num_experts=self.config.n_routed_experts
            + (
                self.config.n_shared_experts
                if rocm_aiter_moe_shared_expert_enabled
                else 0
            ),
302
        )
303
304

        params_dict = dict(self.named_parameters())
305
        loaded_params: set[str] = set()
yangql's avatar
yangql committed
306
307
308
309

        # 判断是否加载"indexer"权重
        model_has_indexer = any("indexer" in param_name for param_name in params_dict.keys())

310
311
312
        for name, loaded_weight in weights:
            if "rotary_emb.inv_freq" in name:
                continue
yangql's avatar
yangql committed
313
314
315
316
317

            #  跳过加载"indexer"权重
            if "indexer" in name and not model_has_indexer:
                logger.info(f"Skipping indexer weight (DSA disabled): {name}")
                continue
318
319
320
            spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
            if spec_layer is None:
                continue
321
322
323
            is_fusion_moe_shared_experts_layer = (
                rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name)
            )
324
            name = self._rewrite_spec_layer_name(spec_layer, name)
325
            for param_name, weight_name, shard_id in stacked_params_mapping:
326
327
328
329
330
331
332
333
334
                # Skip non-stacked layers and experts (experts handled below).
                if weight_name not in name:
                    continue
                # 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.
335
                if ("mlp.experts." in name) and name not in params_dict:
336
                    continue
337
                if is_fusion_moe_shared_experts_layer:
338
                    continue
339
                name_mapped = name.replace(weight_name, param_name)
340
341
342

                # QKV fusion is optional, fall back to normal
                # weight loading if it's not enabled
343
344
345
                if (
                    param_name == "fused_qkv_a_proj"
                ) and name_mapped not in params_dict:
346
                    continue
347
348
                else:
                    name = name_mapped
349

350
351
352
353
354
355
356
357
358
                # Skip loading extra bias for GPTQ models.
                if name.endswith(".bias") and name not in params_dict:
                    continue

                param = params_dict[name]
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                break
            else:
359
360
361
362
363
364
365
366
367
368
369
370
371
372
                # Special handling: when AITER fusion_shared_experts is enabled,
                # checkpoints may provide a single widened shared_experts tensor
                # without explicit expert indices
                # (e.g. ...mlp.shared_experts.gate_proj.weight).
                # For models with multiple shared experts, split that tensor
                # evenly into per-shared-expert slices and load them into
                # appended expert slots mlp.experts.{n_routed_experts + j}.*
                # accordingly.
                num_chunks = 1
                if is_fusion_moe_shared_experts_layer:
                    num_chunks = getattr(self.config, "n_shared_experts", 1) or 1
                    # Determine split axis based on op type
                    # gate/up: ColumnParallel → split along dim 0
                    # down: RowParallel → split along dim 1
373
374
375
376
377
                    split_dim = (
                        1
                        if ("down_proj.weight" in name and loaded_weight.ndim > 1)
                        else 0
                    )
378
379
380
381
                    total = loaded_weight.shape[split_dim]
                    assert total % num_chunks == 0, (
                        f"Shared expert weight dim {total} "
                        f"not divisible by num_chunks {num_chunks}"
382
                    )
383
384
385
386
387
388
389
                    chunk_size = total // num_chunks

                for j in range(num_chunks):
                    chunk_name = name
                    weight_to_load = loaded_weight

                    if is_fusion_moe_shared_experts_layer:
zhuwenwen's avatar
zhuwenwen committed
390
391
392
393
                        if split_dim == 0:
                            weight_to_load = loaded_weight[
                                j * chunk_size : (j + 1) * chunk_size, :
                            ]
394
                        else:
zhuwenwen's avatar
zhuwenwen committed
395
396
397
                            weight_to_load = loaded_weight[
                                :, j * chunk_size : (j + 1) * chunk_size
                            ]
398
399
400
401
402
403
404
405
406
407
                        # Synthesize an expert-style name so expert mapping
                        # can route it
                        chunk_name = name.replace(
                            "mlp.shared_experts",
                            f"mlp.experts.{self.config.n_routed_experts + j}",
                        )

                    # Use expert_params_mapping to locate the destination
                    # param and delegate to its expert-aware weight_loader
                    # with expert_id.
408
                    is_expert_weight = False
409
410
411
412
413
                    for mapping in expert_params_mapping:
                        param_name, weight_name, expert_id, shard_id = mapping
                        if weight_name not in chunk_name:
                            continue

414
415
416
417
                        # Anyway, this is an expert weight and should not be
                        # attempted to load as other weights later
                        is_expert_weight = True

418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
                        # Do not modify `name` since the loop may continue here
                        # Instead, create a new variable
                        name_mapped = chunk_name.replace(weight_name, param_name)

                        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,
                            weight_to_load,
                            name_mapped,
                            shard_id=shard_id,
                            expert_id=expert_id,
                            return_success=True,
                        )
                        if success:
                            if not is_fusion_moe_shared_experts_layer:
                                name = name_mapped
                            else:
                                loaded_params.add(name_mapped)
                            break
                    else:
444
445
446
447
448
449
                        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

450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
                        # Skip loading extra bias for GPTQ models.
                        if name.endswith(".bias") and name not in params_dict:
                            continue

                        name = maybe_remap_kv_scale_name(name, params_dict)
                        if name is None:
                            continue

                        # According to DeepSeek-V3 Technical Report, MTP modules
                        # shares embedding layer. We only load the first weights.
                        if (
                            spec_layer != self.model.mtp_start_layer_idx
                            and ".layers" not in name
                        ):
                            continue

                        param = params_dict[name]
                        weight_loader = getattr(
                            param, "weight_loader", default_weight_loader
                        )
                        weight_loader(param, loaded_weight)
            if not is_fusion_moe_shared_experts_layer:
                loaded_params.add(name)
473
        
zhuwenwen's avatar
zhuwenwen committed
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
        if self.use_llama_nn and self.quant_method is None:
            lay_key_words = [
                "self_attn.eh_proj.weight",
                "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",
                "shared_head.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)
504
505
506
507
508
509
        return loaded_params

    def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
        """
        Rewrite the weight name to match the format of the original model.
        Add .mtp_block for modules in transformer layer block for spec layer
510
        and rename shared layer weights to be top level.
511
512
        """
        spec_layer_weight_names = [
513
514
515
516
517
            "embed_tokens",
            "enorm",
            "hnorm",
            "eh_proj",
            "shared_head",
518
        ]
519
        shared_weight_names = ["embed_tokens"]
520
        spec_layer_weight = False
521
        shared_weight = False
522
523
524
        for weight_name in spec_layer_weight_names:
            if weight_name in name:
                spec_layer_weight = True
525
526
                if weight_name in shared_weight_names:
                    shared_weight = True
527
528
529
                break
        if not spec_layer_weight:
            # treat rest weights as weights for transformer layer block
530
531
532
            name = name.replace(
                f"model.layers.{spec_layer}.", f"model.layers.{spec_layer}.mtp_block."
            )
533
534
535
        elif shared_weight:
            # treat shared weights as top level weights
            name = name.replace(f"model.layers.{spec_layer}.", "model.")
536
        return name