将 python 多重处理与多个大型语言模型配合使用

问：

我正在使用几个LLM（BERT，Word2Vec和Alpa）为文本中的每个单词生成不同的输出。显然，这是相当缓慢的 - 每个单词至少一两秒钟，这对于大型文本来说加起来非常快。因此，我尝试使用多处理来加速此过程，方法是在不同的线程之间拆分工作。

我的理解是，多处理最适用于以下情况：（a） 有很多任务，（b） 每个任务都需要相当长的时间。鉴于这两个原则在我的情况下都是正确的，多处理似乎是减少运行时间的好方法。

我遇到的问题是，酸洗/取消酸洗我的代码被证明在计算上非常昂贵。目前我的代码如下所示：

import os
import torch

import gensim.downloader as vec_api
from transformers import BertTokenizer, BertForMaskedLM, pipeline

from multiprocessing import Pool, cpu_count
import torch.multiprocessing
torch.multiprocessing.set_sharing_strategy('file_system')

vec_model = vec_api.load("fasttext-wiki-news-subwords-300")
tokenizer = BertTokenizer.from_pretrained('bert-large-uncased')
model = BertForMaskedLM.from_pretrained('bert-large-uncased', return_dict=True)
generator = pipeline('text-generation', model="facebook/opt-1.3b")

def pred_word(input):
  # make prediction with BERT using `model`
  # make prediction with Alpa using `generator`
  # get similarities of predictions to original word with Word2Vec using `vec_model`

if __name__ == "main":
  inputs = [] # array of inputs to pass to `pred_word`
  number_of_cores = int(os.getenv('SLURM_CPUS_PER_TASK', cpu_count()))
  chunksize, extra = divmod(len(inputs), 4 * number_of_cores)
  if extra:
    chunksize += 1
  outputs = None
  with Pool(number_of_cores) as pool:
    outputs = pool.map(pred_word, inputs, chunksize=chunksize)

注意：我正在使用 SLURM 来运行此程序。

我的直觉是，在每个线程中加载和/或酸洗 LLM 是该程序花费如此长时间的原因，即使对于非常小的输入大小也是如此。但是，我真的不知道解决这个问题的好方法。