Fix embed scaling + CUDA graphs in Transformers modelling backend (#48010)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Jul 9, 2026 at 00:14 UTC
56da398dac3d86de643f7c12b2bdf68e4cec6d33
1 file changed
+26
-24
vllm/model_executor/models/transformers/base.py
+26
-24
@@ -78,6 +78,17 @@ if TYPE_CHECKING:
78
logger = init_logger(__name__)
79
80
81
+class ScaledVocabParallelEmbedding(VocabParallelEmbedding):
82
+ """`VocabParallelEmbedding` that scales its output."""
83
+
84
+ def __init__(self, *args, embed_scale: float, **kwargs):
85
+ super().__init__(*args, **kwargs)
86
+ self.embed_scale = embed_scale
87
+
88
+ def forward(self, input_: torch.Tensor) -> torch.Tensor:
89
+ return super().forward(input_) * self.embed_scale
90
+
91
+
92
class Base(
93
nn.Module,
94
VllmModel,
@@ -156,22 +167,26 @@ class Base(
167
self.attention_instances = self.create_attention_instances()
168
169
# Input embeddings
159
- self.embed_scale = None
170
input_embeddings = self.model.get_input_embeddings()
171
if not isinstance(input_embeddings, PPMissingLayer):
162
- # Some models scale embeddings inside the input embedding layer
163
- self.embed_scale = getattr(input_embeddings, "embed_scale", None)
172
names = ("embedding_size", "hidden_size")
173
embedding_dim = getattr_iter(self.text_config, names, None)
174
assert embedding_dim is not None
167
- self.model.set_input_embeddings(
168
- VocabParallelEmbedding(
169
- self.text_config.vocab_size,
170
- embedding_dim=embedding_dim,
171
- org_num_embeddings=self.text_config.vocab_size,
172
- quant_config=self.quant_config,
173
- )
175
+ embedding_kwargs = dict(
176
+ num_embeddings=self.text_config.vocab_size,
177
+ embedding_dim=embedding_dim,
178
+ org_num_embeddings=self.text_config.vocab_size,
179
+ quant_config=self.quant_config,
180
)
181
+ embed_scale = getattr(input_embeddings, "embed_scale", None)
182
+ if embed_scale is not None:
183
+ # Some models scale embeddings inside the input embedding layer
184
+ new_input_embeddings = ScaledVocabParallelEmbedding(
185
+ **embedding_kwargs, embed_scale=float(embed_scale)
186
+ )
187
+ else:
188
+ new_input_embeddings = VocabParallelEmbedding(**embedding_kwargs)
189
+ self.model.set_input_embeddings(new_input_embeddings)
190
191
# Initialize any parameters that have not had their modules replaced
192
self.init_parameters(self.model)
@@ -590,10 +605,7 @@ class Base(
605
_init_parameters(module, dtype)
606
607
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
593
- inputs_embeds = self.model.get_input_embeddings()(input_ids)
594
- if self.embed_scale is not None:
595
- inputs_embeds *= self.embed_scale
596
- return inputs_embeds
608
+ return self.model.get_input_embeddings()(input_ids)
609
610
def forward(
611
self,
@@ -608,16 +620,6 @@ class Base(
620
input_ids = None
621
inputs_embeds = intermediate_tensors["hidden_states"]
622
611
- # If the model scales embeddings inside the input embedding layer we must
612
- # ensure they are scaled here since VocabParallelEmbedding will not do it
613
- if (
614
- self.embed_scale is not None
615
- and input_ids is not None
616
- and inputs_embeds is None
617
- ):
618
- inputs_embeds = self.embed_input_ids(input_ids)
619
- input_ids = None
620
-
623
# Add batch dimension before entering Transformers model
624
if input_ids is not None and input_ids.ndim == 1:
625
# [seq_len] -> [1, seq_len]