From 1efed796eeb2555e5194b7a99356100aaeac980e Mon Sep 17 00:00:00 2001 From: tangzhiyi11 Date: Mon, 16 Dec 2024 15:46:27 +0800 Subject: [PATCH] [dlinfer] only compile language_model in vl models (#2893) Co-authored-by: jinminxi104 --- .../backends/dlinfer/ascend/graph_runner.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/lmdeploy/pytorch/backends/dlinfer/ascend/graph_runner.py b/lmdeploy/pytorch/backends/dlinfer/ascend/graph_runner.py index f9664f13ff..e3c5dc4d5e 100644 --- a/lmdeploy/pytorch/backends/dlinfer/ascend/graph_runner.py +++ b/lmdeploy/pytorch/backends/dlinfer/ascend/graph_runner.py @@ -33,10 +33,17 @@ def __init__(self, model: torch.nn.Module, model_config: ModelConfig, dlinfer.graph.config.enable_graph_mode = True self.patch_kernels_custom_op() self.patch_kvcache_static_shape() - self.model = torch.compile(self.model, - fullgraph=True, - dynamic=True, - backend='atbgraph') + if hasattr(self.model, 'language_model'): + self.model.language_model = torch.compile( + self.model.language_model, + fullgraph=True, + dynamic=True, + backend='atbgraph') + else: + self.model = torch.compile(self.model, + fullgraph=True, + dynamic=True, + backend='atbgraph') def check_enable_graph(self): """check enable graph."""