Move Roberta remaining nn.Embedding to VocabParallelEmbedding (#47452)
Signed-off-by: Max de Bayser <mbayser@br.ibm.com>
Max de Bayser committed
Jul 3, 2026 at 05:47 UTC
bbdcbe468636f4174a34ce9396c3e469da045b4b
1 file changed
+3
-5
vllm/model_executor/models/roberta.py
+3
-5
@@ -53,13 +53,11 @@ class RobertaEmbedding(nn.Module):
53
config.vocab_size, config.hidden_size
54
)
55
self.padding_idx = config.pad_token_id
56
- self.position_embeddings = nn.Embedding(
57
- config.max_position_embeddings,
58
- config.hidden_size,
59
- padding_idx=self.padding_idx,
56
+ self.position_embeddings = VocabParallelEmbedding(
57
+ config.max_position_embeddings, config.hidden_size
58
)
59
62
- self.token_type_embeddings = nn.Embedding(
60
+ self.token_type_embeddings = VocabParallelEmbedding(
61
config.type_vocab_size, config.hidden_size
62
)
63
self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)