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

Remove save_pretrained

dev_multi_gpu
duzx16 2023-04-16 21:01:39 +08:00
parent 43b7241e67
commit ffcb1f9394
1 changed files with 2 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
utils.py
View File

@ -47,22 +47,9 @@ def load_model_on_gpus(checkpoint_path: Union[str, os.PathLike], num_gpus: int =
if device_map is None: if device_map is None:
device_map = auto_configure_device_map(num_gpus) device_map = auto_configure_device_map(num_gpus)
try:
model = load_checkpoint_and_dispatch( model = load_checkpoint_and_dispatch(
model, checkpoint_path, device_map=device_map, offload_folder="offload", offload_state_dict=True).half() model, checkpoint_path, device_map=device_map, offload_folder="offload", offload_state_dict=True).half()
except ValueError:
# index.json not found
print(f"index.json not found, auto fixing and saving model to {multi_gpu_model_cache_dir} ...")
assert multi_gpu_model_cache_dir is not None, "using auto fix, cache_dir must not be None"
model.save_pretrained(multi_gpu_model_cache_dir, max_shard_size='2GB')
model = load_checkpoint_and_dispatch(
model, multi_gpu_model_cache_dir, device_map=device_map,
offload_folder="offload", offload_state_dict=True).half()
if tokenizer is not None:
tokenizer.save_pretrained(multi_gpu_model_cache_dir)
print(f"loading model successfully, you should use checkpoint_path={multi_gpu_model_cache_dir} next time")
return model return model