github/DB-GPT
1
0
Fork 0
mirror of https://github.com/csunny/DB-GPT.git synced 2025-08-25 11:29:29 +00:00
Code Issues Packages Projects Releases Wiki Activity

chore: fix shutdown error when not install torch

Browse Source
This commit is contained in:
FangYin Cheng 2023-09-21 12:20:18 +08:00
parent c830598c9e
commit 896af4e16f
2 changed files with 5 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
pilot/model/loader.py
View File

@ -18,6 +18,7 @@ from pilot.logs import logger
def _check_multi_gpu_or_4bit_quantization(model_params: ModelParameters): def _check_multi_gpu_or_4bit_quantization(model_params: ModelParameters):
# TODO: vicuna-v1.5 8-bit quantization info is slow # TODO: vicuna-v1.5 8-bit quantization info is slow
# TODO: support wizardlm quantization, see: https://huggingface.co/WizardLM/WizardLM-13B-V1.2/discussions/5 # TODO: support wizardlm quantization, see: https://huggingface.co/WizardLM/WizardLM-13B-V1.2/discussions/5
# TODO: support internlm quantization
model_name = model_params.model_name.lower() model_name = model_params.model_name.lower()
supported_models = ["llama", "baichuan", "vicuna"] supported_models = ["llama", "baichuan", "vicuna"]
return any(m in model_name for m in supported_models) return any(m in model_name for m in supported_models)

6
pilot/utils/model_utils.py
View File

@ -2,10 +2,12 @@ import logging
def _clear_torch_cache(device="cuda"): def _clear_torch_cache(device="cuda"):
try:
import torch
except ImportError:
return
import gc import gc
import torch
gc.collect() gc.collect()
if device != "cpu": if device != "cpu":
if torch.has_mps: if torch.has_mps: