def load_model_on_gpus(checkpoint_path: Union[str, os.PathLike], num_gpus: int = 2,
device_map: Optional[Dict[str, int]] = None, kwargs) -> Module:
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 dispatch_model
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 = dispatch_model(model, device_map=device_map)
return model
Is there an existing issue for this?
Current Behavior
deepspeed --num_gpus=2 --master_port $MASTER_PORT test02.py \ --deepspeed deepspeed.json \ --do_train \ --train_file /data/train.json \ --test_file /data/test.json \ --prompt_column prompt \ --response_column response \ --history_column history \ --overwrite_cache \ --model_name_or_path ../model \ --output_dir ./output/dataclass-$PROM_TYPE-chatglm2-6b-ft-$LR \ --overwrite_output_dir \ --max_source_length 512 \ --max_target_length 16 \ --per_device_train_batch_size 8 \ --per_device_eval_batch_size 2 \ --gradient_accumulation_steps 2 \ --predict_with_generate \ --max_steps 8000 \ --logging_steps 10 \ --save_steps 4000 \ --learning_rate $LR \ --fp16
from torch.nn import Module from transformers import AutoModel
def auto_configure_device_map(num_gpus: int) -> Dict[str, int]:
transformer.word_embeddings 占用1层
def load_model_on_gpus(checkpoint_path: Union[str, os.PathLike], num_gpus: int = 2, device_map: Optional[Dict[str, int]] = None, kwargs) -> Module: 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 dispatch_model
Anything else?
No response