292 lines
15 KiB
Python
292 lines
15 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
||
from typing import Iterable, Set, Tuple, Optional
|
||
|
||
import torch
|
||
import torch.nn as nn
|
||
from vllm.model_executor.layers.fused_moe import FusedMoE
|
||
|
||
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
|
||
from vllm.forward_context import ForwardContext, get_forward_context
|
||
|
||
from .vars import *
|
||
|
||
from transformers import PretrainedConfig
|
||
|
||
from vllm.config import CacheConfig, ModelConfig, VllmConfig
|
||
from vllm.model_executor.layers.layernorm import RMSNorm
|
||
from vllm.model_executor.layers.quantization import QuantizationConfig
|
||
from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
|
||
from vllm.model_executor.layers.linear import RowParallelLinear
|
||
from vllm.model_executor.models.deepseek_mtp import DeepSeekMultiTokenPredictorLayer as DeepSeekMultiTokenPredictorLayerOrig
|
||
|
||
from vllm.distributed import get_tp_group
|
||
|
||
from vllm.model_executor.models.deepseek_v2 import DeepseekV2DecoderLayer
|
||
|
||
def DeepSeekMultiTokenPredictorLayer__init__(self, vllm_config: VllmConfig, prefix: str) -> None:
|
||
super(DeepSeekMultiTokenPredictorLayerOrig, self).__init__()
|
||
config = vllm_config.model_config.hf_config
|
||
quant_config = vllm_config.quant_config
|
||
self.embed_tokens = VocabParallelEmbedding(
|
||
config.vocab_size,
|
||
config.hidden_size,
|
||
)
|
||
|
||
self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||
self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||
|
||
if USE_PARALLEL_MTP_EH_PROJ:
|
||
self.eh_proj = RowParallelLinear(config.hidden_size * 2,
|
||
config.hidden_size,
|
||
bias=False,
|
||
return_bias=False)
|
||
else:
|
||
self.eh_proj = nn.Linear(config.hidden_size * 2,
|
||
config.hidden_size,
|
||
bias=False)
|
||
|
||
from vllm.model_executor.models.deepseek_mtp import SharedHead
|
||
self.is_v32 = hasattr(config, "index_topk")
|
||
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,
|
||
device="cuda")
|
||
else:
|
||
topk_indices_buffer = None
|
||
self.shared_head = SharedHead(config=config, quant_config=quant_config)
|
||
self.mtp_block = DeepseekV2DecoderLayer(vllm_config, prefix,
|
||
topk_indices_buffer)
|
||
|
||
class DeepSeekMultiTokenPredictorLayer(nn.Module):
|
||
|
||
def forward(
|
||
self,
|
||
input_ids: torch.Tensor,
|
||
positions: torch.Tensor,
|
||
previous_hidden_states: torch.Tensor,
|
||
inputs_embeds: Optional[torch.Tensor] = None,
|
||
spec_step_index: int = 0,
|
||
) -> torch.Tensor:
|
||
|
||
forward_context: ForwardContext = get_forward_context()
|
||
attn_metadata = forward_context.attn_metadata
|
||
if isinstance(attn_metadata, dict):
|
||
attn_metadata = attn_metadata.items().__iter__().__next__()[1]
|
||
|
||
if not hasattr(self, "weight_capture"):
|
||
from vllm_vacc.vllm.model_executor.models.weight_capture.deepseek_weight_capture import DeepseekMTPWegitCapture
|
||
self.weight_capture = DeepseekMTPWegitCapture(self.mtp_block)
|
||
|
||
if inputs_embeds is None:
|
||
inputs_embeds = self.embed_tokens(input_ids)
|
||
assert inputs_embeds is not None
|
||
|
||
if inputs_embeds.shape[0] > 256:
|
||
from vllm_vacc.vllm.model_executor.models.memory.memory_recycling import memory_recycler, DeepseekMTPMemoryRecycler
|
||
deepseek_mtp_layer_input_buffer = None
|
||
if isinstance(memory_recycler, DeepseekMTPMemoryRecycler):
|
||
deepseek_mtp_layer_input_buffer = memory_recycler.DEEPSEEK_MTP_LAYER_INPUT
|
||
|
||
from torch_vacc.vacc.custom_ops import fuse_mtp_stage0
|
||
hidden_states = fuse_mtp_stage0(
|
||
inputs_embeds,
|
||
previous_hidden_states,
|
||
positions,
|
||
self.enorm.weight,
|
||
self.hnorm.weight,
|
||
self.enorm.variance_epsilon,
|
||
world_size=get_tp_group().world_size,
|
||
rank=get_tp_group().rank_in_group,
|
||
group_id=get_tp_group().group_id,
|
||
dev_info=get_tp_group().rank_device_infos,
|
||
output=deepseek_mtp_layer_input_buffer,
|
||
)
|
||
|
||
# if USE_PARALLEL_MTP_EH_PROJ:
|
||
# tp_size = get_tensor_model_parallel_world_size()
|
||
# rank_id = get_tensor_model_parallel_rank()
|
||
# last_dim = hidden_states.shape[-1]
|
||
# if tp_size > 1:
|
||
# hiddens_tp = last_dim//tp_size
|
||
# hidden_states = hidden_states[...,rank_id*hiddens_tp : (rank_id+1)*hiddens_tp]
|
||
|
||
hidden_states = self.eh_proj(hidden_states)
|
||
else:
|
||
hidden_states = torch.vacc.fuse_mtp_allreduce(
|
||
inputs_embeds,
|
||
previous_hidden_states,
|
||
positions,
|
||
self.enorm.weight,
|
||
self.hnorm.weight,
|
||
self.eh_proj.weight,
|
||
self.enorm.variance_epsilon,
|
||
world_size = self.weight_capture.layer_moe.dist_args._0_world_size,
|
||
rank = self.weight_capture.layer_moe.dist_args._1_rank,
|
||
group_id = self.weight_capture.layer_moe.dist_args._2_group_id,
|
||
dev_info = self.weight_capture.layer_moe.dist_args._3_dev_info)
|
||
|
||
if(attn_metadata.prefill_metadata is not None or not USE_DECODER_LAYER_FUSE_MODE):
|
||
hidden_states, residual = self.mtp_block(positions=positions,
|
||
hidden_states=hidden_states,
|
||
residual=None)
|
||
else:
|
||
from torch_vacc.vacc.custom_ops import fuse_mla_moe_v2_allreduce_decode
|
||
layer = self.mtp_block
|
||
layer_id = 0
|
||
|
||
kv_cache = layer.self_attn.mla_attn.kv_cache[forward_context.virtual_engine]
|
||
positions = [p - 1 for p in attn_metadata.decode_metadata.seq_lens]
|
||
cos_cache = [layer.self_attn.mla_attn.impl.rotary_emb.cos_cache[p] for p in positions]
|
||
sin_cache = [layer.self_attn.mla_attn.impl.rotary_emb.sin_cache[p] for p in positions]
|
||
|
||
# 对于MTP Layer来说， residual为None，且需要返回residual
|
||
hidden_states, residual = fuse_mla_moe_v2_allreduce_decode(
|
||
hidden_states = hidden_states,
|
||
residual = None,
|
||
hidden_states_norm_weight = self.weight_capture.layer_moe.attn_args._a_hidden_states_norm_weight[layer_id],
|
||
q_a_proj_weight = self.weight_capture.layer_moe.attn_args._0_merge_q_kv_weights[layer_id],
|
||
q_a_proj_weight_scale_inv = self.weight_capture.layer_moe.attn_args._1_merge_q_kv_scale_inv[layer_id],
|
||
q_a_layernorm_weight = self.weight_capture.layer_moe.attn_args._2_q_a_layernorm_weight[layer_id],
|
||
w_q = self.weight_capture.layer_moe.attn_args._3_W_Q[layer_id],
|
||
w_q_scale = self.weight_capture.layer_moe.attn_args._4_W_Q_scales[layer_id],
|
||
w_uk = self.weight_capture.layer_moe.attn_args._5_W_UK[layer_id],
|
||
w_uk_scale = self.weight_capture.layer_moe.attn_args._6_W_UK_scales[layer_id],
|
||
w_qr = self.weight_capture.layer_moe.attn_args._7_W_QR[layer_id],
|
||
w_qr_scale = self.weight_capture.layer_moe.attn_args._8_W_QR_scales[layer_id],
|
||
kv_a_layernorm_weight = self.weight_capture.layer_moe.attn_args._9_kv_a_layernorm_weight[layer_id],
|
||
sin_cache = sin_cache,
|
||
cos_cache = cos_cache,
|
||
slot_mapping = attn_metadata.slot_mapping,
|
||
kv_cache = kv_cache,
|
||
block_tables = attn_metadata.decode_metadata.block_tables,
|
||
block_group_size = self.weight_capture.layer_moe.attn_args._15_env_blk_grp_size,
|
||
w_uv = self.weight_capture.layer_moe.attn_args._16_W_UV[layer_id],
|
||
w_uv_scale = self.weight_capture.layer_moe.attn_args._17_W_UV_scales[layer_id],
|
||
o_proj_weight = self.weight_capture.layer_moe.attn_args._18_o_proj_weight[layer_id],
|
||
o_proj_weight_scale_inv = self.weight_capture.layer_moe.attn_args._19_o_proj_weight_scale_inv[layer_id],
|
||
# mla params
|
||
seq_lens = attn_metadata.decode_metadata.seq_lens,
|
||
sm_scale = self.weight_capture.layer_moe.attn_args._21_sm_scale,
|
||
head_num = self.weight_capture.layer_moe.attn_args._22_head_num,
|
||
# flash attention
|
||
flash_attention = (USE_FLASH_ATTENTION==1),
|
||
# moe weight
|
||
rms_weight = self.weight_capture.layer_moe.moe_args._0_moe_rms_weight[layer_id],
|
||
mlp_weight_13 = self.weight_capture.layer_moe.moe_args._1_moe_share_mlp_w13[layer_id],
|
||
mlp_weight_2 = self.weight_capture.layer_moe.moe_args._2_moe_share_mlp_w2[layer_id],
|
||
mlp_weight_scale_13 = self.weight_capture.layer_moe.moe_args._3_moe_share_mlp_w13_scale[layer_id],
|
||
mlp_weight_scale_2 = self.weight_capture.layer_moe.moe_args._4_moe_share_mlp_w2_scale[layer_id],
|
||
moe_weight_13 = self.weight_capture.layer_moe.moe_args._5_moe_w13[layer_id],
|
||
moe_weight_2 = self.weight_capture.layer_moe.moe_args._6_moe_w2[layer_id],
|
||
moe_weight_scale_13 = self.weight_capture.layer_moe.moe_args._7_moe_w13_scale[layer_id],
|
||
moe_weight_scale_2 = self.weight_capture.layer_moe.moe_args._8_moe_w2_scale[layer_id],
|
||
mm_weight = self.weight_capture.layer_moe.moe_args._9_gate_weight[layer_id],
|
||
moe_bias = self.weight_capture.layer_moe.moe_args._10_moe_bias[layer_id],
|
||
# moe params
|
||
mlp_block_size_w13 = self.weight_capture.layer_moe.moe_args._11_moe_mlp_w13_block_size,
|
||
mlp_block_size_w2 = self.weight_capture.layer_moe.moe_args._12_moe_mlp_w2_block_size,
|
||
moe_block_size_w13 = self.weight_capture.layer_moe.moe_args._13_moe_w13_block_size,
|
||
moe_block_size_w2 = self.weight_capture.layer_moe.moe_args._14_moe_w2_block_size,
|
||
# vccl info
|
||
world_size = self.weight_capture.layer_moe.dist_args._0_world_size,
|
||
rank = self.weight_capture.layer_moe.dist_args._1_rank,
|
||
group_id = self.weight_capture.layer_moe.dist_args._2_group_id,
|
||
dev_info = self.weight_capture.layer_moe.dist_args._3_dev_info)
|
||
#hidden_states = residual + hidden_states
|
||
hidden_states = residual.add_(hidden_states)
|
||
return hidden_states
|
||
|
||
class DeepSeekMTP(nn.Module):
|
||
|
||
def load_weights(self, weights: Iterable[Tuple[str,
|
||
torch.Tensor]]) -> Set[str]:
|
||
stacked_params_mapping = [
|
||
("gate_up_proj", "gate_proj", 0),
|
||
("gate_up_proj", "up_proj", 1),
|
||
]
|
||
|
||
expert_params_mapping = FusedMoE.make_expert_params_mapping(
|
||
ckpt_gate_proj_name="gate_proj",
|
||
ckpt_down_proj_name="down_proj",
|
||
ckpt_up_proj_name="up_proj",
|
||
num_experts=self.config.n_routed_experts)
|
||
|
||
params_dict = dict(self.named_parameters())
|
||
loaded_params: Set[str] = set()
|
||
from vllm.model_executor.models.deepseek_v2 import get_spec_layer_idx_from_weight_name
|
||
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
|
||
name = self._rewrite_spec_layer_name(spec_layer, name)
|
||
for (param_name, weight_name, shard_id) in stacked_params_mapping:
|
||
# 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.
|
||
if (("mlp.experts." in name) and name not in params_dict):
|
||
continue
|
||
name = name.replace(weight_name, param_name)
|
||
# Skip loading extra bias for GPTQ models.
|
||
if name.endswith(".bias") and name not in params_dict:
|
||
continue
|
||
|
||
param = params_dict[name]
|
||
weight_loader = param.weight_loader
|
||
weight_loader(param, loaded_weight, shard_id)
|
||
break
|
||
else:
|
||
for mapping in expert_params_mapping:
|
||
param_name, weight_name, expert_id, shard_id = mapping
|
||
if weight_name not in name:
|
||
continue
|
||
name = name.replace(weight_name, param_name)
|
||
|
||
param = params_dict[name]
|
||
weight_loader = param.weight_loader
|
||
weight_loader(param,
|
||
loaded_weight,
|
||
name,
|
||
shard_id=shard_id,
|
||
expert_id=expert_id)
|
||
break
|
||
else:
|
||
# Skip loading extra bias for GPTQ models.
|
||
if name.endswith(".bias") and name not in params_dict:
|
||
continue
|
||
|
||
param = params_dict[name]
|
||
weight_loader = getattr(param, "weight_loader",
|
||
default_weight_loader)
|
||
weight_loader(param, loaded_weight)
|
||
loaded_params.add(name)
|
||
|
||
if USE_MERGE_Q_KV_GEN_AND_Q_QR:
|
||
from vllm.model_executor.models.utils import PPMissingLayer
|
||
for layer_id in self.model.layers:
|
||
layer = self.model.layers[layer_id]
|
||
if isinstance(layer_id, PPMissingLayer):
|
||
continue
|
||
layer.mtp_block.self_attn.merge_qkv_weights()
|
||
return loaded_params
|
||
|
||
class SharedHead(nn.Module):
|
||
|
||
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
||
try:
|
||
from torch_vacc.vacc.custom_ops import rms_norm
|
||
return rms_norm(hidden_states, self.norm.weight, output=hidden_states)
|
||
except Exception as e:
|
||
print(f"fuse rms_norm run fail, now use unfused ops: {e}")
|
||
|
||
return self.norm(hidden_states) |