github
/
ChatGLM-6B
mirror of https://github.com/THUDM/ChatGLM-6B
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

Fix multi-gpu loading

pull/518/merge
duzx16 2023-04-17 17:10:09 +08:00
parent 9e61c3ea9b
commit 01e6313abf
1 changed files with 3 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
utils.py
View File

@ -40,16 +40,14 @@ def load_model_on_gpus(checkpoint_path: Union[str, os.PathLike], num_gpus: int =
if num_gpus < 2 and device_map is None:
model = AutoModel.from_pretrained(checkpoint_path, trust_remote_code=True, **kwargs).half().cuda()
else:
from accelerate import load_checkpoint_and_dispatch
from accelerate import dispatch_model
model = AutoModel.from_pretrained(checkpoint_path, trust_remote_code=True, **kwargs)
model = model.eval()
model = AutoModel.from_pretrained(checkpoint_path, trust_remote_code=True, **kwargs).half()
if device_map is None:
device_map = auto_configure_device_map(num_gpus)
model = load_checkpoint_and_dispatch(
model, checkpoint_path, device_map=device_map, offload_folder="offload", offload_state_dict=True).half()
model = dispatch_model(model, device_map=device_map)
return model