Kyungmin Leeandmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
e7a1387e73
Signed-off-by: lkm2835 <lkm2835@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
130 lines
4.4 KiB
Python
130 lines
4.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Inference-only EXAONE-4_5 MTP model."""
|
|
|
|
from collections.abc import Iterable
|
|
|
|
import torch
|
|
from torch import nn
|
|
|
|
from vllm.compilation.decorators import support_torch_compile
|
|
from vllm.config import VllmConfig
|
|
from vllm.logger import init_logger
|
|
from vllm.model_executor.layers.layernorm import RMSNorm
|
|
from vllm.model_executor.layers.linear import ColumnParallelLinear
|
|
from vllm.model_executor.layers.logits_processor import LogitsProcessor
|
|
from vllm.model_executor.layers.vocab_parallel_embedding import (
|
|
ParallelLMHead,
|
|
VocabParallelEmbedding,
|
|
)
|
|
from vllm.model_executor.models.exaone4 import Exaone4DecoderLayer
|
|
from vllm.model_executor.models.exaone_moe_mtp import (
|
|
ExaoneMoeMTP,
|
|
ExaoneMoeMultiTokenPredictor,
|
|
)
|
|
|
|
from .utils import (
|
|
AutoWeightsLoader,
|
|
maybe_prefix,
|
|
)
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
KVCache = tuple[torch.Tensor, torch.Tensor]
|
|
|
|
|
|
@support_torch_compile
|
|
class Exaone4_5MultiTokenPredictor(ExaoneMoeMultiTokenPredictor):
|
|
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
|
nn.Module.__init__(self)
|
|
|
|
model_config = vllm_config.model_config
|
|
quant_config = vllm_config.quant_config
|
|
lora_config = vllm_config.lora_config
|
|
config = model_config.hf_config
|
|
|
|
self.config = config
|
|
lora_vocab = (
|
|
(lora_config.lora_extra_vocab_size * (lora_config.max_loras or 1))
|
|
if lora_config
|
|
else 0
|
|
)
|
|
self.vocab_size = config.vocab_size + lora_vocab
|
|
self.org_vocab_size = config.vocab_size
|
|
|
|
self.mtp_start_layer_idx = config.num_hidden_layers
|
|
self.num_mtp_layers = getattr(config, "num_nextn_predict_layers", 1)
|
|
|
|
self.embed_tokens = VocabParallelEmbedding(
|
|
self.vocab_size,
|
|
config.hidden_size,
|
|
org_num_embeddings=config.vocab_size,
|
|
)
|
|
|
|
self.fc = ColumnParallelLinear(
|
|
self.config.hidden_size * 2,
|
|
self.config.hidden_size,
|
|
gather_output=True,
|
|
bias=False,
|
|
return_bias=False,
|
|
quant_config=quant_config,
|
|
prefix=f"{prefix}.fc",
|
|
)
|
|
self.layers = nn.ModuleList(
|
|
Exaone4DecoderLayer(
|
|
vllm_config.model_config.hf_config,
|
|
quant_config=quant_config,
|
|
prefix=f"{prefix}.layers.{idx}",
|
|
)
|
|
for idx in range(self.num_mtp_layers)
|
|
)
|
|
|
|
self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
|
self.pre_fc_norm_hidden = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
|
self.pre_fc_norm_embedding = RMSNorm(
|
|
config.hidden_size, eps=config.rms_norm_eps
|
|
)
|
|
|
|
|
|
@support_torch_compile
|
|
class Exaone4_5_MTP(ExaoneMoeMTP):
|
|
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
|
config = vllm_config.model_config.hf_config
|
|
self.vllm_config = vllm_config
|
|
self.quant_config = vllm_config.quant_config
|
|
|
|
nn.Module.__init__(self)
|
|
self.config = config
|
|
self.model = Exaone4_5MultiTokenPredictor(
|
|
vllm_config=vllm_config, prefix=maybe_prefix(prefix, "mtp")
|
|
)
|
|
self.unpadded_vocab_size = config.vocab_size
|
|
self.lm_head = ParallelLMHead(
|
|
self.unpadded_vocab_size,
|
|
config.hidden_size,
|
|
org_num_embeddings=config.vocab_size,
|
|
prefix=maybe_prefix(prefix, "lm_head"),
|
|
)
|
|
if config.tie_word_embeddings:
|
|
self.lm_head.weight = self.model.embed_tokens.weight
|
|
self.logits_processor = LogitsProcessor(
|
|
self.unpadded_vocab_size, config.vocab_size
|
|
)
|
|
|
|
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
|
shared_weight_names = ["embed_tokens", "lm_head"]
|
|
|
|
def remap_weight_names(weights):
|
|
for name, weight in weights:
|
|
if name.startswith("mtp."):
|
|
name = name.replace("mtp.", "model.")
|
|
elif any(key in name for key in shared_weight_names):
|
|
if "embed_tokens" in name:
|
|
name = name.replace("language_model.", "")
|
|
else:
|
|
continue
|
|
yield name, weight
|
|
|
|
loader = AutoWeightsLoader(self)
|
|
return loader.load_weights(remap_weight_names(weights))
|