[Help] ptu微调之后,模型加载要5,6分钟有什么办法可以提升速度呢

[1042312930]

Is there an existing issue for this?

• I have searched the existing issues

Current Behavior

加载代码如下
tokenizer = AutoTokenizer.from_pretrained("THUDM\chatglm-6b", trust_remote_code=True)
two_model = AutoModel.from_pretrained("THUDM\chatglm-6b", trust_remote_code=True)
two_prefix_state_dict = torch.load(os.path.join("ptuning/output/adgen-pt-128-2e-2-405sql_train/checkpoint-500", "pytorch_model.bin"))

two_new_prefix_state_dict = {}
for k, v in two_prefix_state_dict.items():
if k.startswith("transformer.prefix_encoder."):
two_new_prefix_state_dict[k[len("transformer.prefix_encoder."):]] = v
two_model.transformer.prefix_encoder.load_state_dict(two_new_prefix_state_dict)

Comment out the following line if you don't use quantization

two_model = two_model.quantize(4)
two_model = two_model.half().cuda()
two_model.transformer.prefix_encoder.float()
two_model = two_model.eval()

Expected Behavior

No response

Steps To Reproduce

训练完成后加载模型(训练的权重也加上)

Environment

- OS:
- Python:3.9
- CUDA Support
- 内存 16G
- GPU 12G

Anything else?

No response