only compile language_model in vl models

InternLM · Dec 13, 2024 · 0d494f8 · 0d494f8
1 parent 422b9f2
commit 0d494f8
Showing 1 changed file with 11 additions and 4 deletions.
diff --git 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."""