GPU Docker环境下运行Optuna+CTGAN脚本出现Segmentation fault错误

问题复现场景 #

执行命令python3 CTGAN_noscale.py --database_name CTGAN_noshift运行CTGAN脚本时触发段错误，faulthandler捕获的完整报错如下：

Fatal Python error: Segmentation fault

Current thread 0x00007f57e97fe700 (most recent call first):
<no Python frame>

Thread 0x00007f593db07740 (most recent call first):
  File "/usr/local/lib/python3.8/dist-packages/torch/autograd/__init__.py", line 145 in backward
  File "/usr/local/lib/python3.8/dist-packages/torch/tensor.py", line 245 in backward
  File "/usr/local/lib/python3.8/dist-packages/ctgan/synthesizers/ctgan.py", line 374 in fit
  File "CTGAN_noscale.py", line 140 in objective
  File "CTGAN_noscale.py", line 162 in <lambda>
  File "/usr/local/lib/python3.8/dist-packages/optuna/_optimize.py", line 216 in _run_trial
  File "/usr/local/lib/python3.8/dist-packages/optuna/_optimize.py", line 162 in _optimize_sequential
  File "/usr/local/lib/python3.8/dist-packages/optuna/_optimize.py", line 65 in _optimize
  File "/usr/local/lib/python3.8/dist-packages/optuna/study.py", line 401 in optimize
  File "CTGAN_noscale.py", line 162 in run_CTGAN
  File "CTGAN_noscale.py", line 210 in <module>
Segmentation fault (core dumped)

异常特征 #

• 相同脚本在原服务器运行完全正常，更换服务器后才触发崩溃
• 仅在启用GPU的Docker容器内运行时复现问题，裸机运行、不启用GPU的Docker容器内运行均无异常

根因定位 #

该段错误和Optuna无关，从调用栈可以看到崩溃触发在PyTorch反向传播阶段，结合仅GPU Docker容器复现的特征，属于CUDA环境兼容性问题：
新服务器宿主机的NVIDIA驱动版本、Docker容器内安装的CUDA toolkit版本、PyTorch版本三者不匹配，GPU计算时触发底层非法内存访问。部分旧版本CTGAN依赖的PyTorch版本和新架构GPU（如30/40系消费级显卡、A10/A100等计算卡）也存在兼容性问题，开启GPU加速后会触发段错误。

修复方案 #

• 校验版本匹配关系：分别在宿主机和容器内执行nvidia-smi，确认容器内CUDA版本不高于宿主机NVIDIA驱动支持的最高CUDA版本
• 统一依赖版本：将容器内的PyTorch、torchvision、torchaudio版本调整为和正常运行的原服务器完全一致，优先使用PyTorch官方匹配对应CUDA版本的安装命令
• 启动容器时添加--ipc=host参数，避免Docker容器共享内存不足导致的PyTorch GPU计算崩溃
• 临时添加环境变量CUDA_LAUNCH_BLOCKING=1运行脚本，可获取更精准的CUDA报错信息，进一步定位具体不兼容的算子

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