加载检查点分片需要太长时间

Question

我对生成式人工智能非常陌生。我有 64GB RAM 和 20GB GPU。我使用了 Huggingface 的一些开源模型，并使用 Python 简单地用开箱即用的模型提示并显示结果。我使用

save_pretrained

将模型下载到本地，然后尝试从本地加载模型。有用。但每次运行 python 文件都需要 10 多分钟才能显示结果。

有一个步骤

Loading checkpoint shards

每次需要6-7分钟。我做错了什么吗？为什么它每次都必须加载一些东西，即使模型是从本地引用的。

我尝试使用

local_files_only=True, cache_dir=cache_dir, low_cpu_mem_usage=True, max_shard_size="200MB"

，但没有解决时间问题。

如何在用户可用的情况下直接提示已保存的模型，而不会造成太大的延迟。任何帮助将不胜感激

Answer 1

我遇到了完全相同的问题，我通过在使用

safe_serialization=True

方法时设置

save_pretrained()

来修复它。希望这对你有用。但是，我确实想知道加载普通 .bin 格式的模型时发生了什么。

Answer 2

问题源于

from_pretrained

方法在每次脚本运行时重新加载检查点分片。

为了避免这种情况，请考虑在 Jupyter 笔记本中运行代码，这允许您使用

from_pretrained

在单独的单元中加载模型一次，并保持内核运行。可以在其他单元中使用该模型来生成输出。

设置方法如下：

# Cell 1: Load the model (Run this once)
from transformers import YourModelClass
model = YourModelClass.from_pretrained("path/to/local/model")

# Cell 2: Use the model (Run this as needed)
output = model.generate_some_output()
print(output)