[Fix]: support deepseek-vl2-tiny model (#5552)

Co-authored-by: bppps <zouyu.zzx@alibaba-inc.com>
This commit is contained in:
ZXN
2025-04-26 17:52:53 +08:00
committed by GitHub
co-authored by bppps
parent feda9b11b3
commit 04d0123fd9
6 changed files with 80 additions and 6 deletions
+12 -2
View File
@@ -382,8 +382,14 @@ class DeepseekModel(nn.Module):
input_ids: torch.Tensor,
positions: torch.Tensor,
forward_batch: ForwardBatch,
input_embeds: torch.Tensor = None,
) -> torch.Tensor:
hidden_states = self.embed_tokens(input_ids)
if input_embeds is None:
hidden_states = self.embed_tokens(input_ids)
else:
hidden_states = input_embeds
residual = None
for i in range(len(self.layers)):
layer = self.layers[i]
@@ -416,14 +422,18 @@ class DeepseekForCausalLM(nn.Module):
)
self.logits_processor = LogitsProcessor(config)
def get_input_embeddings(self) -> nn.Embedding:
return self.model.embed_tokens
@torch.no_grad()
def forward(
self,
input_ids: torch.Tensor,
positions: torch.Tensor,
forward_batch: ForwardBatch,
input_embeds: torch.Tensor = None,
) -> torch.Tensor:
hidden_states = self.model(input_ids, positions, forward_batch)
hidden_states = self.model(input_ids, positions, forward_batch, input_embeds)
return self.logits_processor(
input_ids, hidden_states, self.lm_head, forward_batch
)