最近做实验遇到了一个非常棘手的随机崩溃问题，实在摸不着头脑，想请各位大佬帮忙分析下：

我用PyTorch 2.6.0 + Optuna 结合CUDA 12.4训练Transformer模型时，频繁随机触发段错误（Segmentation Fault）。之前完全相同的训练代码单独运行时一切正常，唯独在通过Optuna的study.optimize()调度训练时才会出问题——哪怕设置n_jobs=1也躲不过这个坑。

崩溃的位置非常深，直接卡在Transformer的前馈块里，而且没有任何Python异常抛出，进程直接就挂掉了，报错信息如下：

Fatal Python error: Segmentation fault
File ".../torch/nn/modules/transformer.py", line 947, in _ff_block
File ".../torch/nn/functional.py", line 1704, in relu

最头疼的是这个错误还不是确定性的，有时候跑几次trial都没事，有时候第一次trial就直接崩，完全摸不准规律。

可复现的简化代码 #

import torch
import torch.nn as nn
import optuna

device = "cuda:0"

# --- 简易Transformer模型 ---
class TinyTransformer(nn.Module):
    def __init__(self, d_model=128, nhead=2, num_layers=2):
        super().__init__()
        layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=nhead, batch_first=True
        )
        self.encoder = nn.TransformerEncoder(layer, num_layers=num_layers)
        self.proj = nn.Linear(d_model, 1)

    def forward(self, x):
        return self.proj(self.encoder(x))

def objective(trial):
    model = TinyTransformer().to(device)
    opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
    loss_fn = nn.MSELoss()
    
    x = torch.randn(8, 48, 128, device=device)
    y = torch.randn(8, 48, 1, device=device)
    
    for _ in range(10):
        opt.zero_grad(set_to_none=True)
        out = model(x)
        loss = loss_fn(out, y)
        loss.backward()
        opt.step()
    return float(loss)

if __name__ == "__main__":
    study = optuna.create_study(direction="minimize")
    # 注意：只有调用这个方法时才会触发崩溃
    study.optimize(objective, n_trials=5, n_jobs=1)

我的疑问 #

• Optuna和PyTorch 2.6在CUDA环境下是否存在已知的兼容性问题？比如进程fork或者线程处理机制上的冲突？
• TransformerEncoderLayer默认启用的Flash Attention会不会是引发这个随机崩溃的元凶？