deepseek_mtp.py 20.3 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
zhuwenwen's avatar
zhuwenwen committed
40
from vllm import _custom_ops as ops
zhuwenwen's avatar
zhuwenwen committed
41
from vllm.model_executor.layers.quantization.blockwise_int8 import BlockInt8Config
42
import vllm.envs as envs
43

44
logger = init_logger(__name__)
45
46
47
48
49
50


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

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


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

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

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

79
        self.device = current_platform.device_type
80

81
        self.is_v32 = hasattr(config, "index_topk")
82
83
84
85
86
87
        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,
88
                device=self.device,
89
            )
90
91
        else:
            topk_indices_buffer = None
92

93
94
95
96
        self.shared_head = SharedHead(
            config=config, prefix=prefix, quant_config=quant_config
        )
        self.mtp_block = DeepseekV2DecoderLayer(
97
98
99
100
            vllm_config,
            prefix,
            config=self.config,
            topk_indices_buffer=topk_indices_buffer,
101
        )
102
103
104
105
106
107

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        previous_hidden_states: torch.Tensor,
108
        inputs_embeds: torch.Tensor | None = None,
109
110
111
112
        spec_step_index: int = 0,
    ) -> torch.Tensor:
        assert inputs_embeds is not None
        # masking inputs at position 0, as not needed by MTP
113
114
115
116
117
118
119
120
        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)
121

122
123
124
            hidden_states = self.eh_proj(
                torch.cat([inputs_embeds, previous_hidden_states], dim=-1)
            )
125

126
127
128
        hidden_states, residual = self.mtp_block(
            positions=positions, hidden_states=hidden_states, residual=None
        )
129
        hidden_states = residual + hidden_states
130
        return hidden_states
131
132
133
134
135
136
137
138
139


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
140

141
142
143
144
145
146
147
148
149
150
151
        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,
                )
            }
        )
152
153
154
        self.embed_tokens = VocabParallelEmbedding(
            config.vocab_size,
            config.hidden_size,
155
            prefix=maybe_prefix(prefix, "embed_tokens"),
156
        )
157
158
        self.logits_processor = LogitsProcessor(config.vocab_size)

159
    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
160
161
        return self.embed_tokens(input_ids)

162
163
164
165
166
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        previous_hidden_states: torch.Tensor,
167
        inputs_embeds: torch.Tensor | None = None,
168
169
        spec_step_idx: int = 0,
    ) -> torch.Tensor:
170
171
        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)
172
        current_step_idx = spec_step_idx % self.num_mtp_layers
173
        return self.layers[str(self.mtp_start_layer_idx + current_step_idx)](
174
175
176
177
            input_ids,
            positions,
            previous_hidden_states,
            inputs_embeds,
178
            current_step_idx,
179
180
181
182
183
184
185
        )

    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        spec_step_idx: int = 0,
    ) -> torch.Tensor:
186
187
188
189
190
        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)
        )
191
192
193
        return logits


194
@support_torch_compile
195
class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
196
197
198
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        super().__init__()
        self.config = vllm_config.model_config.hf_config
199
200
201
        self.model = DeepSeekMultiTokenPredictor(
            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
        )
202
203
        # Set MoE hyperparameters
        self.set_moe_parameters()
zhuwenwen's avatar
zhuwenwen committed
204
205
206
207
208
209
210
        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
211
            # The AWQ layer of MTP uses BlockInt8W8A8.
yangql's avatar
yangql committed
212
            if self.quant_method == "moe_wna16" or self.quant_method == "awq_marlin":
zhuwenwen's avatar
zhuwenwen committed
213
                vllm_config.quant_config = BlockInt8Config(is_checkpoint_int8_serialized=True, weight_block_size=[128,128])
zhuwenwen's avatar
zhuwenwen committed
214
        self.use_llama_nn = os.environ.get('LLAMA_NN') == '1'
215
                
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233

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

235
236
    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.model.embed_input_ids(input_ids)
237
238
239

    def forward(
        self,
240
        input_ids: torch.Tensor | None,
241
        positions: torch.Tensor,
242
        hidden_states: torch.Tensor,
243
244
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
245
246
        spec_step_idx: int = 0,
    ) -> torch.Tensor:
247
248
249
        hidden_states = self.model(
            input_ids, positions, hidden_states, inputs_embeds, spec_step_idx
        )
250
251
252
253
254
255
        return hidden_states

    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        spec_step_idx: int = 0,
256
    ) -> torch.Tensor | None:
257
        return self.model.compute_logits(hidden_states, spec_step_idx)
258

259
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
260
261
262
        rocm_aiter_moe_shared_expert_enabled = (
            rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
        )
263
264
265
        stacked_params_mapping = [
            ("gate_up_proj", "gate_proj", 0),
            ("gate_up_proj", "up_proj", 1),
266
267
            ("fused_qkv_a_proj", "q_a_proj", 0),
            ("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1),
268
269
        ]

270
        expert_params_mapping = SharedFusedMoE.make_expert_params_mapping(
271
            self,
272
273
274
            ckpt_gate_proj_name="gate_proj",
            ckpt_down_proj_name="down_proj",
            ckpt_up_proj_name="up_proj",
275
276
277
278
279
280
            num_experts=self.config.n_routed_experts
            + (
                self.config.n_shared_experts
                if rocm_aiter_moe_shared_expert_enabled
                else 0
            ),
281
        )
282
283

        params_dict = dict(self.named_parameters())
284
        loaded_params: set[str] = set()
285
286
287
288
289
290
        for name, loaded_weight in weights:
            if "rotary_emb.inv_freq" in name:
                continue
            spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
            if spec_layer is None:
                continue
291
292
293
            is_fusion_moe_shared_experts_layer = (
                rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name)
            )
294
            name = self._rewrite_spec_layer_name(spec_layer, name)
295
            for param_name, weight_name, shard_id in stacked_params_mapping:
296
297
298
299
300
301
302
303
304
                # 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.
305
                if ("mlp.experts." in name) and name not in params_dict:
306
                    continue
307
                if is_fusion_moe_shared_experts_layer:
308
                    continue
309
                name_mapped = name.replace(weight_name, param_name)
310
311
312

                # QKV fusion is optional, fall back to normal
                # weight loading if it's not enabled
313
314
315
                if (
                    param_name == "fused_qkv_a_proj"
                ) and name_mapped not in params_dict:
316
                    continue
317
318
                else:
                    name = name_mapped
319

320
321
322
323
324
325
326
327
328
                # 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:
329
330
331
332
333
334
335
336
337
338
339
340
341
342
                # 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
343
344
345
346
347
                    split_dim = (
                        1
                        if ("down_proj.weight" in name and loaded_weight.ndim > 1)
                        else 0
                    )
348
349
350
351
                    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}"
352
                    )
353
354
355
356
357
358
359
                    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:
360
361
362
363
364
                        chunk_slice = slice(j * chunk_size, (j + 1) * chunk_size)
                        if loaded_weight.ndim == 1:
                            weight_to_load = loaded_weight[chunk_slice]
                        elif split_dim == 0:
                            weight_to_load = loaded_weight[chunk_slice, :]
365
                        else:
366
                            weight_to_load = loaded_weight[:, chunk_slice]
367
368
369
370
371
372
373
374
375
376
                        # 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.
377
                    is_expert_weight = False
378
379
380
381
382
                    for mapping in expert_params_mapping:
                        param_name, weight_name, expert_id, shard_id = mapping
                        if weight_name not in chunk_name:
                            continue

383
384
385
386
                        # Anyway, this is an expert weight and should not be
                        # attempted to load as other weights later
                        is_expert_weight = True

387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
                        # 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:
413
414
415
416
417
418
                        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

419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
                        # 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)
442
        
zhuwenwen's avatar
zhuwenwen committed
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
        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)
473
474
475
476
477
478
        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
479
        and rename shared layer weights to be top level.
480
481
        """
        spec_layer_weight_names = [
482
483
484
485
486
            "embed_tokens",
            "enorm",
            "hnorm",
            "eh_proj",
            "shared_head",
487
        ]
488
        shared_weight_names = ["embed_tokens"]
489
        spec_layer_weight = False
490
        shared_weight = False
491
492
493
        for weight_name in spec_layer_weight_names:
            if weight_name in name:
                spec_layer_weight = True
494
495
                if weight_name in shared_weight_names:
                    shared_weight = True
496
497
498
                break
        if not spec_layer_weight:
            # treat rest weights as weights for transformer layer block
499
500
501
            name = name.replace(
                f"model.layers.{spec_layer}.", f"model.layers.{spec_layer}.mtp_block."
            )
502
503
504
        elif shared_weight:
            # treat shared weights as top level weights
            name = name.replace(f"model.layers.{spec_layer}.", "model.")
505
        return name