biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[V1] Refactor model executable interface for all text-only language models (#10374)

Browse Source 
Signed-off-by: Roger Wang <ywang@roblox.com>
This commit is contained in:
Roger Wang
2024-11-16 21:18:46 -08:00
committed by GitHub
parent 4fd9375028
commit 643ecf7b11
43 changed files with 483 additions and 90 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
vllm/model_executor/models/gpt2.py
View File

@@ -209,6 +209,9 @@ class GPT2Model(nn.Module):
make_empty_intermediate_tensors_factory(["hidden_states"],
config.n_embd))
def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
return self.wte(input_ids)
def forward(
self,
input_ids: torch.Tensor,
@@ -220,7 +223,7 @@ class GPT2Model(nn.Module):
) -> Union[torch.Tensor, IntermediateTensors]:
if get_pp_group().is_first_rank:
if inputs_embeds is None:
inputs_embeds = self.wte(input_ids)
inputs_embeds = self.get_input_embeddings(input_ids)
position_embeds = self.wpe(position_ids)
hidden_states = inputs_embeds + position_embeds
else:
@@ -262,7 +265,7 @@ class GPT2LMHeadModel(nn.Module, SupportsPP):
self.transformer.make_empty_intermediate_tensors)
def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
return self.transformer.wte(input_ids)
return self.transformer.get_input_embeddings(input_ids)
def forward(
self,