From c313af0639b9d1c27a8d5e65f73b6ed683e63ee3 Mon Sep 17 00:00:00 2001 From: saber <3082548039@qq.com> Date: Mon, 27 Mar 2023 22:20:23 +0800 Subject: [PATCH] Add support for manual device_map; auto-config if None --- utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/utils.py b/utils.py index 2bf13c2..0c0d757 100644 --- a/utils.py +++ b/utils.py @@ -38,13 +38,15 @@ def auto_configure_device_map(num_gpus: int) -> Dict[str, int]: def load_model_on_gpus(checkpoint_path: Union[str, os.PathLike], num_gpus: int = 2, multi_gpu_model_cache_dir: Union[str, os.PathLike] = "./temp_model_dir", + device_map: Optional[Dict[str, int]] = None, tokenizer: Optional[PreTrainedTokenizer] = None, **kwargs) -> Module: from accelerate import load_checkpoint_and_dispatch model = AutoModel.from_pretrained(checkpoint_path, trust_remote_code=True, **kwargs) model = model.eval() - device_map = auto_configure_device_map(num_gpus) + if device_map is None: + device_map = auto_configure_device_map(num_gpus) try: model = load_checkpoint_and_dispatch( model, checkpoint_path, device_map=device_map, offload_folder="offload", offload_state_dict=True).half()