Skip to content

Commit

Permalink
update warning info
Browse files Browse the repository at this point in the history
  • Loading branch information
tangzhiyi11 committed Nov 14, 2024
1 parent b96b7dd commit c354cb3
Showing 1 changed file with 2 additions and 6 deletions.
8 changes: 2 additions & 6 deletions lmdeploy/pytorch/backends/dlinfer/ascend/graph_runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,6 @@ def __init__(self, model: torch.nn.Module, model_config: ModelConfig,
super().__init__(model, model_config, cache_config, backend_config,
device)

self.supported_model = ['Llama3-8B', 'Llama2-7B', 'Qwen2-7B']
self.enable_graph = self.check_enable_graph()
if self.enable_graph:
import dlinfer.graph
Expand All @@ -43,11 +42,8 @@ def check_enable_graph(self):
warnings.warn(
'\n\n'
'**********************************************************\n'
' The following models were tested in graph mode of\n'
" device_type 'ascend' when tp=1:\n"
f" {', '.join(self.supported_model)}\n"
' Other LLaMa-like models may work in graph mode, please\n'
' check the result yourself!\n'
' Graph mode is an experimental feature.We are now support\n'
' both dense and moe model with bf16/fp16 datatype.\n'
' If graph mode does not work correctly with your model,\n'
' please use eager mode instead.\n'
'**********************************************************\n\n',
Expand Down

0 comments on commit c354cb3

Please sign in to comment.