本地加载HuggingFace模型时加载分片耗时过长求解

针对你遇到的本地缓存模型但每次加载分片耗时7-10分钟的问题，以下是几个实用的解决办法：

1. 将分片模型合并为单文件（一次性操作） #

你可以先加载一次模型（尽管这次慢，但只需要做一次），然后重新保存为单文件格式，后续加载就不会再有分片加载的问题：

from transformers import AutoModelForCausalLM, AutoTokenizer

# 首次加载模型（耗时较长，仅执行一次）
model = AutoModelForCausalLM.from_pretrained('./cache/model', low_cpu_mem_usage=True)
tokenizer = AutoTokenizer.from_pretrained('./cache/model')

# 保存为单文件格式（safe_serialization=False 生成单文件而非分片）
model.save_pretrained('./cache/model_single_file', safe_serialization=False)
tokenizer.save_pretrained('./cache/model_single_file')

之后每次加载模型时，直接使用新的单文件路径：

model = AutoModelForCausalLM.from_pretrained('./cache/model_single_file', low_cpu_mem_usage=True)
tokenizer = AutoTokenizer.from_pretrained('./cache/model_single_file')

2. 优化加载参数（无需修改模型文件） #

如果不想重新保存模型，可以在加载时添加以下参数提升速度：

• low_cpu_mem_usage=True：优化加载时的CPU内存占用，避免不必要的内存复制，减少加载时间。
• device_map="auto"：让transformers自动将模型层分配到GPU，减少CPU与GPU之间的数据传输耗时，尤其适合大模型。

示例代码：

model = AutoModelForCausalLM.from_pretrained(
    './cache/model',
    low_cpu_mem_usage=True,
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained('./cache/model')

3. 量化加载模型（兼顾速度与性能） #

如果你的业务场景允许，可以使用4bit/8bit量化加载模型，模型体积会大幅缩小，加载速度显著提升，同时对推理性能影响很小：

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained(
    './cache/model',
    load_in_4bit=True,  # 或 load_in_8bit=True
    device_map="auto",
    low_cpu_mem_usage=True
)
tokenizer = AutoTokenizer.from_pretrained('./cache/model')

补充说明 #

分片模型加载慢的核心原因是需要逐个读取多个shard文件并在内存中拼接，涉及多次磁盘IO操作；合并为单文件后只需一次磁盘读取，加载速度会大幅提升。你提到的safe_serialization参数，在保存模型时设置为False就会生成单文件，所以通过上述“先加载再重新保存”的方式就能应用这个方案。

内容的提问来源于stack exchange，提问作者SJay_747