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

Merge pull request #1 from smartmark-pro/fix-not-use-gpu 

ptuning或者微调加载到GPU设备
pull/334/head
Mark 2023-07-19 12:04:22 +08:00 committed by GitHub
parent 738e015cac 83abf3ca38
commit dd4bb36244
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
ptuning/main.py
View File

@ -112,7 +112,7 @@ def main():
if model_args.ptuning_checkpoint is not None:
# Evaluation
# Loading extra state dict of prefix encoder
model = AutoModel.from_pretrained(model_args.model_name_or_path, config=config, trust_remote_code=True)
model = AutoModel.from_pretrained(model_args.model_name_or_path, config=config, trust_remote_code=True, device='cuda')
prefix_state_dict = torch.load(os.path.join(model_args.ptuning_checkpoint, "pytorch_model.bin"))
new_prefix_state_dict = {}
for k, v in prefix_state_dict.items():
@ -120,7 +120,7 @@ def main():
new_prefix_state_dict[k[len("transformer.prefix_encoder."):]] = v
model.transformer.prefix_encoder.load_state_dict(new_prefix_state_dict)
else:
model = AutoModel.from_pretrained(model_args.model_name_or_path, config=config, trust_remote_code=True)
model = AutoModel.from_pretrained(model_args.model_name_or_path, config=config, trust_remote_code=True, device='cuda')
if model_args.quantization_bit is not None:
print(f"Quantized to {model_args.quantization_bit} bit")