使用Optuna+Transformers时n_jobs>1触发Tensor.item()元张量错误

问题背景 #

当使用Optuna为HuggingFace语言模型微调任务调参，将n_jobs设为大于1的并行运行模式时，会触发RuntimeError: Tensor.item() cannot be called on meta tensors错误。特殊表现包括：

• 交互式会话中首次并行运行时，一个trial成功另一个失败；后续同会话内所有运行都会失败
• 并行失败后即使改回n_jobs=1，仍会报错，直到重启Python解释器
• 仅返回纯张量的目标函数可正常并行运行

原因分析 #

1. 多进程环境下的状态污染：Optuna默认使用多进程（multiprocessing）并行，HuggingFace模型加载时的设备/张量状态可能在进程间意外共享，导致部分进程中的模型参数被加载到meta设备（懒加载相关逻辑），后续调用.item()时触发错误。
2. 交互式会话的残留影响：失败进程的meta张量状态会污染主进程的PyTorch上下文，即使后续切换回单进程模式，仍会复用错误的张量状态。
3. 张量跨进程序列化问题：直接返回PyTorch张量而非Python数值时，跨进程序列化可能导致张量状态异常。

解决方案 #

1. 目标函数内完全独立初始化资源并显式指定设备 #

每个trial进程独立加载模型、tokenizer，并将所有张量和模型移动到实际设备（CPU/GPU），避免meta设备残留：

import torch
import optuna
from transformers import AutoModelForCausalLM, AutoTokenizer

def objective(trial):
    # 每个进程独立加载模型，禁用自动设备映射
    model = AutoModelForCausalLM.from_pretrained('gpt2', device_map=None)
    tokenizer = AutoTokenizer.from_pretrained('gpt2')
    
    # 显式指定设备，确保模型和张量在实际设备上
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model.to(device)
    
    # 输入张量同步到对应设备
    inputs = tokenizer('This is a test sentence.', return_tensors='pt').to(device)
    inputs['labels'] = inputs['input_ids']
    
    outputs = model(**inputs)
    # 将loss转换为Python数值后返回，避免跨进程序列化问题
    return outputs.loss.item()

study = optuna.create_study()
study.optimize(objective, n_trials=2, n_jobs=2)
print(f'Best param is {study.best_params} with value {study.best_value}')

2. 清理交互式会话残留状态 #

如果在Jupyter/IPython等交互式环境中运行，出现错误后必须重启内核；命令行环境则重启Python解释器，彻底清除进程残留的异常张量状态。

3. 可选：切换Optuna并行后端 #

若多进程模式仍有问题，可尝试使用线程后端（需注意GIL对CPU密集型任务的性能影响）：

# 使用SQLite存储避免线程间状态冲突
study = optuna.create_study(storage="sqlite:///optuna_study.db", load_if_exists=True)
study.optimize(objective, n_trials=2, n_jobs=2, backend="threading")

额外优化建议 #

• 对于大型模型，可使用torch.multiprocessing的共享内存机制复用模型权重，但需确保初始化逻辑完全隔离，避免meta设备问题。
• GPU环境下，可设置torch.cuda.set_per_process_memory_fraction(0.4, device=device)限制单进程内存占比，避免OOM。

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