環境變數

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")