From 27ae6763647920119dcd50d147ca365fdc138945 Mon Sep 17 00:00:00 2001 From: Kyungmin Lee <30465912+lkm2835@users.noreply.github.com> Date: Mon, 11 May 2026 19:25:31 +0900 Subject: [PATCH] Fix EXAONE-4.5 to align with Transformers update (#42246) Signed-off-by: lkm2835 --- vllm/model_executor/models/exaone4_5.py | 4 -- vllm/model_executor/models/exaone4_5_mtp.py | 62 ++++++++++++++++++--- 2 files changed, 53 insertions(+), 13 deletions(-) diff --git a/vllm/model_executor/models/exaone4_5.py b/vllm/model_executor/models/exaone4_5.py index 1eac43ccb0..b44708466c 100644 --- a/vllm/model_executor/models/exaone4_5.py +++ b/vllm/model_executor/models/exaone4_5.py @@ -23,7 +23,6 @@ import torch import torch.nn as nn from transformers.models.exaone4_5 import ( Exaone4_5_Config, - Exaone4_5_ImageProcessor, Exaone4_5_Processor, ) from transformers.models.exaone4_5.configuration_exaone4_5 import Exaone4_5_VisionConfig @@ -304,9 +303,6 @@ class Exaone4_5_ProcessingInfo(Qwen2VLProcessingInfo): **kwargs, ) - def get_image_processor(self, **kwargs: object) -> Exaone4_5_ImageProcessor: - return Exaone4_5_ImageProcessor(**kwargs) - @MULTIMODAL_REGISTRY.register_processor( Exaone4_5_MultiModalProcessor, diff --git a/vllm/model_executor/models/exaone4_5_mtp.py b/vllm/model_executor/models/exaone4_5_mtp.py index 99bf724bda..7711f72e42 100644 --- a/vllm/model_executor/models/exaone4_5_mtp.py +++ b/vllm/model_executor/models/exaone4_5_mtp.py @@ -9,6 +9,7 @@ from torch import nn from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig +from vllm.distributed.parallel_state import get_pp_group from vllm.logger import init_logger from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ColumnParallelLinear @@ -22,6 +23,7 @@ from vllm.model_executor.models.exaone_moe_mtp import ( ExaoneMoeMTP, ExaoneMoeMultiTokenPredictor, ) +from vllm.sequence import IntermediateTensors from .interfaces import ( MultiModalEmbeddings, @@ -48,6 +50,7 @@ class Exaone4_5MultiTokenPredictor(ExaoneMoeMultiTokenPredictor): quant_config = vllm_config.quant_config lora_config = vllm_config.lora_config config = model_config.hf_config + text_config = config.text_config self.config = config lora_vocab = ( @@ -58,18 +61,18 @@ class Exaone4_5MultiTokenPredictor(ExaoneMoeMultiTokenPredictor): 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.mtp_start_layer_idx = text_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, + text_config.hidden_size, org_num_embeddings=config.vocab_size, ) self.fc = ColumnParallelLinear( - self.config.hidden_size * 2, - self.config.hidden_size, + text_config.hidden_size * 2, + text_config.hidden_size, gather_output=True, bias=False, return_bias=False, @@ -78,27 +81,68 @@ class Exaone4_5MultiTokenPredictor(ExaoneMoeMultiTokenPredictor): ) self.layers = nn.ModuleList( Exaone4DecoderLayer( - vllm_config.model_config.hf_config, + text_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.norm = RMSNorm(text_config.hidden_size, eps=text_config.rms_norm_eps) + self.pre_fc_norm_hidden = RMSNorm( + text_config.hidden_size, eps=text_config.rms_norm_eps + ) self.pre_fc_norm_embedding = RMSNorm( - config.hidden_size, eps=config.rms_norm_eps + text_config.hidden_size, eps=text_config.rms_norm_eps ) def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.embed_tokens(input_ids) + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + hidden_states: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + spec_step_idx: int = 0, + ) -> torch.Tensor: + if get_pp_group().is_first_rank: + if inputs_embeds is None: + inputs_embeds = self.get_input_embeddings(input_ids) + assert hidden_states.shape[-1] == inputs_embeds.shape[-1] + inputs_embeds = self.pre_fc_norm_embedding(inputs_embeds) + hidden_states = self.pre_fc_norm_hidden(hidden_states) + hidden_states = torch.cat([inputs_embeds, hidden_states], dim=-1) + hidden_states = self.fc(hidden_states) + residual = None + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + residual = intermediate_tensors["residual"] + + current_step_idx = spec_step_idx % self.num_mtp_layers + hidden_states, residual = self.layers[current_step_idx]( + positions=positions, + hidden_states=hidden_states, + residual=residual, + ) + + if not get_pp_group().is_last_rank: + return IntermediateTensors( + {"hidden_states": hidden_states, "residual": residual} + ) + + hidden_states = self.norm(hidden_states) + return hidden_states + @support_torch_compile class Exaone4_5_MTP(ExaoneMoeMTP, SupportsMultiModal): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): config = vllm_config.model_config.hf_config + text_config = config.text_config self.vllm_config = vllm_config self.quant_config = vllm_config.quant_config @@ -110,7 +154,7 @@ class Exaone4_5_MTP(ExaoneMoeMTP, SupportsMultiModal): self.unpadded_vocab_size = config.vocab_size self.lm_head = ParallelLMHead( self.unpadded_vocab_size, - config.hidden_size, + text_config.hidden_size, org_num_embeddings=config.vocab_size, prefix=maybe_prefix(prefix, "lm_head"), )