[Bugfix] Transformers backend: apply learned lm_head.bias for tied-embedding models (#46835)
Signed-off-by: John Langford <jl@hunch.net> Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
John committed
Jun 29, 2026 at 04:59 UTC
0e207dac784e6b217b8dc1f44ae3985b1f216b50
3 files changed
+26
-15
vllm/model_executor/layers/vocab_parallel_embedding.py
+7
-10
@@ -542,19 +542,16 @@ class ParallelLMHead(VocabParallelEmbedding):
542
)
543
self.quant_config = quant_config
544
if bias:
545
- self.bias = Parameter(
546
- torch.empty(self.num_embeddings_per_partition, dtype=params_dtype)
547
- )
548
- set_weight_attrs(
549
- self.bias,
550
- {
551
- "output_dim": 0,
552
- "weight_loader": self.weight_loader,
553
- },
554
- )
545
+ self._register_bias()
546
else:
547
self.register_parameter("bias", None)
548
549
+ def _register_bias(self):
550
+ data = torch.empty(self.num_embeddings_per_partition, dtype=self.params_dtype)
551
+ self.bias = Parameter(data, requires_grad=False)
552
+ weight_attrs = dict(output_dim=0, weight_loader=self.weight_loader)
553
+ set_weight_attrs(weight=self.bias, weight_attrs=weight_attrs)
554
+
555
def tie_weights(self, embed_tokens: VocabParallelEmbedding):
556
"""Tie the weights with word embeddings."""
557
return self.quant_method.tie_weights(self, embed_tokens)
vllm/model_executor/models/transformers/base.py
+1
-4
@@ -640,10 +640,7 @@ class Base(
640
return hidden_states, aux_hidden_states
641
return hidden_states
642
643
- def load_weights(
644
- self,
645
- weights: Iterable[tuple[str, torch.Tensor]],
646
- ) -> set[str]:
643
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
644
loader = AutoWeightsLoader(
645
self,
646
skip_prefixes=self.skip_prefixes,
vllm/model_executor/models/transformers/causal.py
+18
-1
@@ -16,6 +16,7 @@
16
# limitations under the License.
17
"""Transformers modeling backend mixin for causal language models."""
18
19
+from collections.abc import Iterable
20
from typing import TYPE_CHECKING
21
22
from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -61,6 +62,22 @@ class CausalMixin(VllmModelForTextGeneration):
62
else:
63
self.lm_head = PPMissingLayer()
64
65
+ def load_weights(self, weights: Iterable[tuple[str, "torch.Tensor"]]) -> set[str]:
66
+ """A thin wrapper around `Base.load_weights` to handle the lm_head bias."""
67
+
68
+ lm_head_bias = set()
69
+
70
+ def auto_load_lm_head_bias(weights):
71
+ for name, weight in weights:
72
+ if name.endswith("lm_head.bias") and self.pp_group.is_last_rank:
73
+ self.lm_head._register_bias()
74
+ self.lm_head.bias.weight_loader(self.lm_head.bias, weight)
75
+ lm_head_bias.add(name)
76
+ else:
77
+ yield name, weight
78
+
79
+ return super().load_weights(auto_load_lm_head_bias(weights)) | lm_head_bias
80
+
81
def compute_logits(self, hidden_states: "torch.Tensor") -> "torch.Tensor | None":
65
- logits = self.logits_processor(self.lm_head, hidden_states)
82
+ logits = self.logits_processor(self.lm_head, hidden_states, self.lm_head.bias)
83
return logits