环境变量

HF_ENABLE_PARALLEL_LOADING

默认情况下，此功能处于禁用状态。启用基于 Torch 和 Safetensor 的权重并行加载。可以显著减少加载大型模型所需的时间，通常能将速度提升约 50%。

可设置为等于 "false" 或 "true" 的字符串。例如：os.environ["HF_ENABLE_PARALLEL_LOADING"] = "true"。

例如：在 AWS EC2 g4dn.metal 实例上，启用此功能后，加载 facebook/opt-30b 大约只需 30 秒，而禁用此功能则需要约 55 秒。

在使用此环境变量之前请先进行性能分析，因为它不会对较小的模型产生加速效果。

import os

os.environ["HF_ENABLE_PARALLEL_LOADING"] = "true"

from transformers import pipeline

model = pipeline(task="text-generation", model="facebook/opt-30b", device_map="auto")

HF_PARALLEL_LOADING_WORKERS

确定启用并行加载时应使用的线程数。默认值为 8。

如果正在加载的文件数量少于指定的线程数，则实际生成的线程数将等于文件数量。

例如：如果指定 8 个工作线程，但只有 2 个文件，则只会生成 2 个工作线程。

根据需要进行调整。

import os

os.environ["HF_ENABLE_PARALLEL_LOADING"] = "true"
os.environ["HF_PARALLEL_LOADING_WORKERS"] = "4"

from transformers import pipeline

model = pipeline(task="text-generation", model="facebook/opt-30b", device_map="auto")