如何在单GPU上运行多Optuna试验以充分利用显存？

我有一个用于超参数调优的Optuna脚本，目前通过以下命令在4块GPU上运行：

CUDA_VISIBLE_DEVICES=0 python 'optuna script.py' &
CUDA_VISIBLE_DEVICES=1 python 'optuna script.py' &
CUDA_VISIBLE_DEVICES=2 python 'optuna script.py' &
CUDA_VISIBLE_DEVICES=3 python 'optuna script.py' &

每块GPU拥有24GB显存，但当前每块GPU仅同时运行1个试验（每个试验仅占用约1GB显存，算力消耗极低），导致调优过程非常缓慢。我希望通过在每块GPU上同时运行多个试验，充分利用可用显存。

以下是我的代码相关部分：

def objective(trial):
    # Defined my neural network architecture, data loaders, loss function, and optimizer here
    ...

def main():
    storage_url = "mysql://optuna_user:optuna_password@localhost/optuna_db"
    study_name = "distributed-relapse-prediction-full"
    study = optuna.create_study(study_name=study_name, direction="maximize", storage=storage_url, load_if_exists=True)
    study.optimize(objective, n_trials=100)

if __name__ == "__main__":
    main()

请问是否可以修改脚本或命令，实现每块GPU运行多个试验，以高效利用全部可用显存？恳请提供相关指导或建议。

方法1：单GPU进程内启用多试验并行 #

Optuna原生支持在单个进程内并行运行多个试验，只需修改study.optimize()的n_jobs参数，指定单GPU上同时运行的试验数量。按你的显存条件，每块GPU可并行20+个试验，修改代码如下：

def main():
    storage_url = "mysql://optuna_user:optuna_password@localhost/optuna_db"
    study_name = "distributed-relapse-prediction-full"
    study = optuna.create_study(study_name=study_name, direction="maximize", storage=storage_url, load_if_exists=True)
    # 设置n_jobs为单GPU并行的试验数，示例设为20
    study.optimize(objective, n_trials=100, n_jobs=20)

启动命令保持原有形式即可，每个GPU对应一个进程，进程内部自动并行多个试验：

CUDA_VISIBLE_DEVICES=0 python 'optuna script.py' &
CUDA_VISIBLE_DEVICES=1 python 'optuna script.py' &
CUDA_VISIBLE_DEVICES=2 python 'optuna script.py' &
CUDA_VISIBLE_DEVICES=3 python 'optuna script.py' &

额外注意：如果objective函数涉及CUDA操作，需确保每个试验的模型绑定到当前GPU，可在函数开头添加：

import torch
torch.cuda.set_device(0)  # 进程已被限制到单个GPU，此处固定为0即可

方法2：单GPU启动多个独立Optuna进程 #

若不想修改代码，可直接在每个GPU上启动多个独立进程，比如每块GPU启动20个：

# GPU 0启动20个进程
for i in {1..20}; do
    CUDA_VISIBLE_DEVICES=0 python 'optuna script.py' &
done
# GPU 1启动20个进程
for i in {1..20}; do
    CUDA_VISIBLE_DEVICES=1 python 'optuna script.py' &
done
# GPU 2启动20个进程
for i in {1..20}; do
    CUDA_VISIBLE_DEVICES=2 python 'optuna script.py' &
done
# GPU 3启动20个进程
for i in {1..20}; do
    CUDA_VISIBLE_DEVICES=3 python 'optuna script.py' &
done

这种方式无需改代码，但进程数量较多，需注意不要超出系统进程上限。

通用注意事项 #

• 先小批量测试：先在单GPU上跑5个并行试验，确认显存占用和稳定性后，再逐步增加并行数。
• 数据库性能：多个试验同时更新MySQL存储，需确保数据库能承受并发写入，避免锁等待拖慢调优。
• 显存泄漏排查：若长时间运行后显存持续上升，需在objective函数末尾添加torch.cuda.empty_cache()，释放未使用的显存资源。

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