[BUG] <title>Loading checkpoint shards takes too long

[sunyclj]

是否已有关于该错误的issue或讨论？ | Is there an existing issue / discussion for this?

• 我已经搜索过已有的issues和讨论 | I have searched the existing issues / discussions

该问题是否在FAQ中有解答？ | Is there an existing answer for this in FAQ?

• 我已经搜索过FAQ | I have searched FAQ

当前行为 | Current Behavior

Loading checkpoint shards takes too long,about 12-15 minutes

What is the reason?

期望行为 | Expected Behavior

No response

复现方法 | Steps To Reproduce

import torch
from PIL import Image
from transformers import AutoModel, AutoTokenizer
import time

torch.manual_seed(0)

model = AutoModel.from_pretrained('/home/jovyan/work/suny/ocr/MiniCPM-V-2_6', trust_remote_code=True,attn_implementation='sdpa', torch_dtype=torch.bfloat16) # sdpa or flash_attention_2, no eager
model = model.eval().cuda()
tokenizer = AutoTokenizer.from_pretrained('/home/jovyan/work/suny/ocr/MiniCPM-V-2_6', trust_remote_code=True)

image = Image.open('./21301.jpg').convert('RGB')

First round chat

question = 'XXX'

msgs = [{'role': 'user', 'content': [image,question]}]

answer = model.chat(
image=None,
msgs=msgs,
tokenizer=tokenizer
)
print(answer)

运行环境 | Environment

- OS:ubuntu 20.04
- Python:python 3.10
- Transformers:4.40.0
- PyTorch:2.1.2
- CUDA (`python -c 'import torch; print(torch.version.cuda)'`):12.3
- GPU: A100*1

备注 | Anything else?

No response