[torch.compile] Adding torch compile to vision-language models (#9946)

This commit is contained in:
Yongzao
2024-11-03 03:56:05 +08:00
committed by GitHub
parent 1b73ab2a1f
commit ae5279a163
3 changed files with 21 additions and 8 deletions

View File

@@ -564,8 +564,13 @@ class MiniCPMVBaseModel(nn.Module, SupportsMultiModal, SupportsPP):
vlm_embeddings, _ = self.get_embedding(input_ids, image_inputs)
# always pass the input via `inputs_embeds`
# to make sure the computation graph is consistent
# for `torch.compile` integration
input_ids = None
output = self.llm(
input_ids=None,
input_ids=input_ids,
positions=positions,
kv_caches=kv_caches,
attn_metadata=attn_metadata,