github/DB-GPT
1
0
Fork 0
mirror of https://github.com/csunny/DB-GPT.git synced 2025-07-23 20:26:15 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix load model gpu oom

Browse Source
This commit is contained in:
csunny 2023-04-29 23:02:13 +08:00
parent ca29dacc37
commit eca14bc038
1 changed files with 5 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
pilot/model/loader.py
View File

@ -7,6 +7,8 @@ from transformers import (
AutoModelForCausalLM, AutoModelForCausalLM,
) )
from fastchat.serve.compression import compress_module
class ModerLoader: class ModerLoader:
kwargs = {} kwargs = {}
@ -29,6 +31,9 @@ class ModerLoader:
if debug: if debug:
print(model) print(model)
if load_8bit:
compress_module(model, self.device)
# if self.device == "cuda": # if self.device == "cuda":
# model.to(self.device) # model.to(self.device)