JobLib运行GPT-2并行推理出现PicklingError序列化报错问题问询

JobLib默认使用pickle模块序列化需要传递给子进程的任务及所有关联对象，你的代码存在两个导致序列化失败的核心问题：

• 全局定义的HuggingFace pipeline 实例generator包含大量PyTorch张量、模型权重、C/C++扩展绑定对象，这类对象本身不支持pickle序列化
• query函数依赖全局的generator对象，JobLib在打包任务时会尝试序列化generator，直接触发PicklingError

优先推荐批量推理方案，无需多进程就能大幅提升效率，仅需加载一次模型，不会额外占用显存：

import time
from transformers import pipeline, set_seed
set_seed(42)
generator = pipeline('text-generation', model='gpt2')

invoc_list = ['his her criminal charges include but are not limited to:', 'the defendant was sentenced to:', 'his her education history:']
entries = 4

%%time
# 直接传入批量输入，pipeline内部自动并行处理
batch_result = generator(invoc_list, min_length=25, max_length=50, num_return_sequences=entries)
# 按输入分组处理输出
my_list = []
for idx in range(len(invoc_list)):
    current_res = [d['generated_text'].split(invoc_list[idx])[1].strip() for d in batch_result[idx*entries : (idx+1)*entries]]
    my_list.append(current_res)

如果确实需要使用多进程实现，需要将模型初始化逻辑迁移到子进程内，避免在主进程创建模型后传递给子进程：

import time
from joblib import Parallel, delayed
from transformers import pipeline, set_seed

invoc_list = ['his her criminal charges include but are not limited to:', 'the defendant was sentenced to:', 'his her education history:']
entries = 4

def query(payload, entries):
    # 模型在子进程内部初始化，不需要序列化传递
    set_seed(42)
    generator = pipeline('text-generation', model='gpt2')
    list_dict = generator(payload, min_length=25, max_length=50, num_return_sequences=entries)
    output = [d['generated_text'].split(payload)[1].strip() for d in list_dict]
    return output

%%time
# 注意n_jobs不要超过你的显存上限，每个进程会单独加载一份模型
my_list = Parallel(n_jobs=2)(delayed(query)(invoc, entries) for invoc in invoc_list)

注意：多进程方案每个子进程都会加载一份GPT-2模型，单份GPT-2占约500MB显存，需要根据你的显存容量调整n_jobs数量，避免显存溢出。

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